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 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)
boxes = boxes.astype(np.int32)
img = cv2.imread(des_img_path)
img = DrawRect(img, boxes)
cv2.imwrite(save_path, img)
num += 1
print(save_path)
if __name__ == "__main__":
img_path = "../../SeaShips/JPEGImages/001223.jpg"
save_src_path = "../per_data/001223.jpg"
save_des_path = "../per_data/001224.jpg"
anno_path = "../../SeaShips/Annotations/001223.xml"
data = SeashipDataset("../../SeaShips", None)
#img, anno = data.getitem2(1222)
#boxes = anno['boxes']
boxes = readxml(anno_path)
print(boxes)
img = cv2.imread(img_path)
cv2.imwrite(save_src_path, img)
img, boxes = horizontal_flip(img, boxes)
#img = DrawRect(img, boxes)
cv2.imwrite(save_des_path, img)
'''
img, boxes = rotate(img, boxes, 15)
optimizer,
data_loader,
device,
epoch,
print_freq=10)
lr_scheduler.step()
evaluate(model, data_loader_test, device=device)
print("That's it!")
if __name__ == '__main__':
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
num_classes = 7
dataset = SeashipDataset('../', get_transform(train=True))
dataset_test = SeashipDataset('../', get_transform(train=False))
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=True,
num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=4,
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])])
'''