def get_model(num_classes, device="cuda:0"):
# 加载backbone时不需要加载官方模型
backbone = resnet_fpn_backbone('resnet50',
pretrained=False,
trainable_layers=3)
model = FasterRCNN(backbone, num_classes=91) # 加载官方模型，这里num_classes不能修改
# 加载fasterrcnn_resnet50_fpn官方模型
# https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
weights_dict = torch.load(
"C:/Users/lixiao/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
map_location=device)
missing_keys, unexpected_keys = model.load_state_dict(weights_dict,
strict=False)
overwrite_eps(model, 0.0)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes) # 修改predictor
return model
def test_train():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.load_state_dict(torch.load("./mask_rcnn_2.pth"))
model.to(device)
model.train()
data = SeashipDataset("../SeaShips", None)
imgs = []
targets = []
for i in range(2):
img, target = data.__getitem__(i)
imgs.append(F.to_tensor(img).to(device))
target = {k: v.to(device) for k, v in target.items()}
targets.append(target)
targets.append(target)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
result = model.forward(imgs, targets)
def test():
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.load_state_dict(torch.load("../model/mask_rcnn_5_2_002.pth"))
model.to(device)
model.eval()
data = SeashipDataset("../../SeaShips", None)
evaluator = Evaluator(7, 0.8)
with open("./test_data", "r") as f:
lines = f.readlines()
test_list = [int(line) for line in lines]
print(len(test_list))
result_list = []
target_list = []
batch_size = 2
for idx in range(0, len(test_list), batch_size):
imgs = []
targets = []
try:
for i in range(idx, idx + batch_size):
img, target = data.__getitem__(test_list[i] - 1)
imgs.append(F.to_tensor(img).to(device))
target = {k: v.to(device) for k, v in target.items()}
targets.append(target)
results = model.forward(imgs)
#result_list.extend(results)
#target_list.extend(targets)
for result, target in zip(results, targets):
target = {k: v.cpu().detach() for k, v in target.items()}
result = {k: v.cpu().detach() for k, v in result.items()}
result_list.append(result)
target_list.append(target)
#evaluator.get_label(result, target)
except:
print(str(traceback.format_exc()))
if idx % 12 == 0:
print(idx)
evaluator.get_ap(result_list, target_list)
def res_mask():
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.load_state_dict(torch.load("../model/mask_rcnn_10_2_005.pth"))
model.to(device)
model.eval()
data = SeashipDataset("../SeaShips", None)
with open("./train_data", "r") as f:
lines = f.readlines()
train_list = [int(line) for line in lines]
print(len(train_list))
batch_size = 2
for idx in range(0, len(train_list), batch_size):
imgs = []
for i in range(idx, idx + batch_size):
img, target = data.__getitem__(train_list[i])
imgs.append(F.to_tensor(img).to(device))
original_image_sizes = [img.shape[-2:] for img in imgs]
result = model.forward(imgs)
for j, res in enumerate(result):
scores = res['scores'].cpu().detach().numpy()
masks = res["masks"].cpu()
index = np.where(scores > 0.9) #只要9分以上的
masks = masks[index]
masks = torch.where(masks > 0.5, torch.full_like(masks, 1),
torch.full_like(masks, 0))
m = torch.zeros(original_image_sizes[0])
for mask in masks:
m += mask[0]
m = torch.where(m > 0.5, torch.full_like(m, 1),
torch.full_like(m, 0))
img_mask = TensorToPIL(m)
mask_name = "{:0>6d}.png".format(idx + j + 1)
path = os.path.join("./res_mask/", mask_name)
img_mask.save(path)
print(path)
def get_model_detection(model_name,
weight_loss,
backbone,
pretrained,
pretrained_backbone,
original,
detection_score_thres,
use_non_player,
use_soft_nms,
anchor_sizes=[32, 64, 128, 256, 512],
n_channel_backbone=5,
use_context=False,
nms_thres=0.4,
use_track_branch=False):
if use_non_player:
num_classes = 3
else:
num_classes = 2
if model_name == 'frcnn_fpn':
if backbone == 'resnet18':
model = fasterrcnn_resnet18_fpn(
weight_loss=weight_loss,
pretrained_backbone=pretrained_backbone,
num_classes=2,
detection_score_thres=detection_score_thres,
anchor_sizes=anchor_sizes,
n_channel_backbone=n_channel_backbone,
box_nms_thresh=nms_thres,
use_soft_nms=use_soft_nms,
use_context=use_context,
use_track_branch=use_track_branch)
if backbone == 'resnet50':
model = fasterrcnn_resnet50_fpn(
weight_loss=weight_loss,
pretrained=pretrained,
pretrained_backbone=pretrained_backbone,
num_classes=2,
anchor_sizes=anchor_sizes,
detection_score_thres=detection_score_thres,
n_channel_backbone=n_channel_backbone,
box_nms_thresh=nms_thres,
use_soft_nms=use_soft_nms,
use_context=use_context,
use_track_branch=use_track_branch)
if not original:
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
else:
'model not available'
return model
def get_model(num_classes, device="cuda:0"):
# # https://download.pytorch.org/models/vgg16-397923af.pth
# # 如果使用mobilenetv2的话就下载对应预训练权重并注释下面三行，接着把mobilenetv2模型对应的两行代码注释取消掉
# vgg_feature = vgg(model_name="vgg16", weights_path="./backbone/vgg16.pth").features
# backbone = torch.nn.Sequential(*list(vgg_feature._modules.values())[:-1]) # 删除feature中最后的maxpool层
# backbone.out_channels = 512
# # https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
# # backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features
# # backbone.out_channels = 1280 # 设置对应backbone输出特征矩阵的channels
# anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
# aspect_ratios=((0.5, 1.0, 2.0),))
# roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], # 在哪些特征层上进行roi pooling
# output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
# sampling_ratio=2) # 采样率
# model = FasterRCNN(backbone=backbone,
# num_classes=num_classes,
# rpn_anchor_generator=anchor_generator,
# box_roi_pool=roi_pooler)
# 加载backbone时不需要加载官方模型
backbone = resnet_fpn_backbone('resnet50',
pretrained=False,
trainable_layers=3)
model = FasterRCNN(backbone, num_classes=91) # 加载官方模型，这里num_classes不能修改
# 加载fasterrcnn_resnet50_fpn官方模型
# https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
weights_dict = torch.load(
"C:/Users/lixiao/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
map_location=device)
missing_keys, unexpected_keys = model.load_state_dict(weights_dict,
strict=False)
overwrite_eps(model, 0.0)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes) # 修改predictor
return model
def test_one(idx):
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.load_state_dict(torch.load("../model/mask_rcnn_5_2_002.pth"))
model.to(device)
model.eval()
data = SeashipDataset("../../SeaShips", None)
evaluator = Evaluator(7, 0.5)
targets = []
image, target = data.__getitem__(idx)
img_var = [F.to_tensor(image).to(device)]
original_image_sizes = [img.shape[-2:] for img in img_var]
targets.append(target)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
result = model.forward(img_var)
#print(target['labels'])
#print(result)
evaluator.get_label(result[0], target)
boxes = result[0]['boxes'].cpu().detach().numpy()
scores = result[0]['scores'].cpu().detach().numpy()
masks = result[0]['masks'].cpu()
index = np.where(scores > 0.5)
boxes = boxes[index]
masks = masks[index]
masks = torch.where(masks > 0.5, torch.full_like(masks, 1),
torch.full_like(masks, 0))
m = torch.zeros(original_image_sizes[0])
for mask in masks:
m += mask[0]
m = torch.where(m > 0.5, torch.full_like(m, 1), torch.full_like(m, 0))
image = image.convert("RGBA")
mask = TensorToPIL(m)
mask.convert("RGBA")
drawRectangle(image, boxes)
image = drawMasks(image, mask)
res_name = "../result/test{:0>6d}.png".format(idx)
image.save(res_name)
def weakly_supervision_train():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.load_state_dict(
torch.load("../model/mask_rcnn_5_2_002.pth")) #加载已训练的模型
model.to(device)
model.open_weakly_supervision_train() #打开弱监督训练方式
for name, params in model.named_parameters():
if 'mask' not in name: #冻结mask分支以外的参数
params.requires_grad = False
optimizer = torch.optim.SGD(model.parameters(),
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
num_epochs = 5
batch_size = 2
loss_sum = 0
data = SeashipDataset("../../SeaShips", None)
with open("./train_data", "r") as f:
lines = f.readlines()
train_list = [int(line) for line in lines]
print(len(train_list))
for epoch in range(num_epochs):
for idx in range(0, len(train_list), batch_size):
try:
imgs = []
targets = []
for i in range(idx, idx + batch_size):
img, target = data.getitem2(train_list[i] - 1, epoch)
imgs.append(F.to_tensor(img).to(device))
target = {k: v.to(device) for k, v in target.items()}
targets.append(target)
original_image_sizes = [img.shape[-2:] for img in imgs]
loss_dict, result = model.forward(imgs, targets)
losses = sum(loss for loss in loss_dict.values())
loss_sum += losses
#print(result)
optimizer.zero_grad()
losses.backward()
optimizer.step()
for j, res in enumerate(result):
scores = res['scores'].cpu().detach().numpy()
masks = res["masks"].cpu()
#print(masks[0].shape)
index = np.where(scores > 0.9) #只要9分以上的
masks = masks[index]
masks = torch.where(masks > 0.5, torch.full_like(masks, 1),
torch.full_like(masks, 0))
m = torch.zeros(original_image_sizes[0])
for mask in masks:
m += mask[0]
m = torch.where(m > 0.5, torch.full_like(m, 1),
torch.full_like(m, 0))
img_mask = TensorToPIL(m)
data.updatemask(idx + j, img_mask, epoch)
except:
logger.error(str(traceback.format_exc()))
if idx % 10 == 0:
#print("[%d]rpn_loss: %f" %(idx, loss_sum))
logger.debug("[%d]total_loss: %f" % (idx, loss_sum))
loss_sum = 0
torch.save(model.state_dict(), "./mask_rcnn_weakly_5_2_002.pth")
def train_mask():
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.to(device)
model.train()
data = SeashipDataset("../../SeaShips", None)
optimizer = torch.optim.SGD(model.parameters(),
lr=0.002,
momentum=0.9,
weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 5
batch_size = 2
loss_sum = 0
loss_classifier = 0
loss_box_reg = 0
loss_mask = 0
with open("./train_data", "r") as f:
lines = f.readlines()
train_list = [int(line) for line in lines]
print(len(train_list))
for epoch in range(num_epochs):
lr_scheduler = None
if epoch == 0:
warmup_factor = 1. / 1000
warmup_iters = min(1000, 5000 - 1)
lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters,
warmup_factor)
for idx in range(0, len(train_list), batch_size):
try:
imgs = []
targets = []
for i in range(idx, idx + batch_size):
img, target = data.__getitem__(train_list[i] - 1)
imgs.append(F.to_tensor(img).to(device))
target = {k: v.to(device) for k, v in target.items()}
targets.append(target)
loss_dict = model.forward(imgs, targets)
losses = sum(loss for loss in loss_dict.values())
loss_sum += losses
#loss_classifier += loss_dict['loss_classifier'].values()
#loss_box_reg += loss_dict['loss_box_reg'].values()
#loss_mask += loss_dict['loss_mask'].values()
optimizer.zero_grad()
losses.backward()
optimizer.step()
if lr_scheduler is not None:
lr_scheduler.step()
except:
logger.error(str(traceback.format_exc()))
if idx % 12 == 0:
logger.debug("[%d]total_loss: %f" % (idx, loss_sum))
#logger.debug("[%d]loss: %f loss_classifier: %f loss_box_reg: %f loss_mask: %f" %(idx, loss_sum, loss_classifier, loss_box_reg, loss_mask))
loss_sum = 0
#loss_classifier = 0
#loss_box_reg = 0
#loss_mask = 0
torch.save(model.state_dict(), "../model/mask_rcnn_5_2_002.pth")
logger.debug("train successfully!")\
def test_mask():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
model = maskrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 7)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer, 7)
model.load_state_dict(torch.load("./mask_rcnn_2.pth"))
model.to(device)
#model.train()
model.eval()
data = SeashipDataset("../SeaShips", None)
targets = []
image, target = data.__getitem__(1888)
img_var = [F.to_tensor(image).to(device)]
original_image_sizes = [img.shape[-2:] for img in img_var]
targets.append(target)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
#result = model.forward(img_var, targets)
result = model.forward(img_var)
'''
batch_size = 2
label_struct = dict()
for idx in range(5000,7000,batch_size):
imgs = []
targets = []
for i in range(idx, idx+batch_size):
img, target = data.__getitem__(i)
imgs.append(F.to_tensor(img).to(device))
target = {k : v.to(device) for k, v in target.items()}
targets.append(target)
result = model.forward(imgs)
for res in result:
bbox
'''
print(result)
print(target['labels'])
boxes = result[0]['boxes'].cpu().detach().numpy()
scores = result[0]['scores'].cpu().detach().numpy()
masks = result[0]['masks'].cpu()
index = np.where(scores > 0.9)
boxes = boxes[index]
masks = masks[index]
masks = torch.where(masks > 0.5, torch.full_like(masks, 1),
torch.full_like(masks, 0))
m = torch.zeros(original_image_sizes[0])
print(m.shape)
for mask in masks:
m += mask[0]
m = torch.where(m > 0.5, torch.full_like(m, 1), torch.full_like(m, 0))
img_mask = TensorToPIL(m)
img_mask.save("./result/res11.png")
#masks[pos] = 1
#print(boxes)
#print(masks[0][0][int(boxes[0][1]):int(boxes[0][3]),int(boxes[0][0]):int(boxes[0][2])])
'''
def get_model_detection(model_name, weight_loss, backbone_name, pretrained,
pretrained_backbone, original, detection_score_thres, use_soft_nms,
anchor_sizes=[32, 64, 128, 256, 512], n_channel_backbone=5, use_context=False,
nms_thres=0.4, use_track_branch=False):
num_classes = 2
if model_name == 'frcnn_fpn':
if backbone_name == 'resnet8':
model = fasterrcnn_resnet8_fpn(weight_loss=weight_loss,
pretrained_backbone=pretrained_backbone, num_classes=2,
detection_score_thres=detection_score_thres, anchor_sizes=anchor_sizes,
n_channel_backbone=n_channel_backbone, use_soft_nms=use_soft_nms,
use_context=use_context)
if backbone_name == 'resnet18':
model = fasterrcnn_resnet18_fpn(weight_loss=weight_loss,
pretrained_backbone=pretrained_backbone, num_classes=2,
detection_score_thres=detection_score_thres, anchor_sizes=anchor_sizes,
n_channel_backbone=n_channel_backbone, use_soft_nms=use_soft_nms,
nms_thres=nms_thres,
use_context=use_context,
use_track_branch=use_track_branch)
if backbone_name == 'resnet34':
model = fasterrcnn_resnet34_fpn(weight_loss=weight_loss,
pretrained_backbone=pretrained_backbone, num_classes=2,
detection_score_thres=detection_score_thres, anchor_sizes=anchor_sizes,
n_channel_backbone=n_channel_backbone, use_soft_nms=use_soft_nms,
use_context=use_context)
if backbone_name == 'resnet50':
model = fasterrcnn_resnet50_fpn(weight_loss=weight_loss,
pretrained=pretrained, pretrained_backbone=pretrained_backbone,
num_classes=2, anchor_sizes=anchor_sizes,
detection_score_thres=detection_score_thres,
n_channel_backbone=n_channel_backbone, use_soft_nms=use_soft_nms,
nms_thres=nms_thres,
use_context=use_context, use_track_branch=use_track_branch)
if backbone_name == 'detnet59':
model = fasterrcnn_detnet59_fpn(weight_loss=weight_loss,
pretrained=pretrained, pretrained_backbone=pretrained_backbone,
num_classes=2, anchor_sizes=anchor_sizes,
detection_score_thres=detection_score_thres,
use_soft_nms=use_soft_nms, nms_thres=nms_thres, use_context=use_context)
if not original:
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
elif model_name == 'rpn':
if backbone_name == 'resnet18':
backbone_model = resnet_fpn_backbone('resnet18', pretrained_backbone)
model = RPN(backbone_model, weight_loss=weight_loss)
if backbone_name == 'resnet50':
backbone_model = resnet_fpn_backbone('resnet50', pretrained_backbone)
model = RPN(backbone_model, weight_loss=weight_loss)
if pretrained:
load_model = fasterrcnn_resnet50_fpn(pretrained=True)
model.backbone = load_model.backbone
model.rpn = load_model.rpn
else:
'model not available'
return model