def __init__(self, kind="COCO", config=None, model_dir=None):
self.config = config
self.model_dir = model_dir
self.previous_preloaded_model = self.find_last()
if kind == "COCO":
pretrained_model_path = "./mask_rcnn_coco.pth"
cfg = CocoConfig()
cfg.NAME = config.NAME
else:
raise NotImplementedError("Provide a valid pretrained model name")
super().__init__(cfg, model_dir)
self.load_weights(pretrained_model_path)
# change classifier
self.classifier = Classifier(256, config.POOL_SIZE, config.IMAGE_SHAPE,
config.NUM_CLASSES)
# change mask
self.mask = Mask(256, config.MASK_POOL_SIZE, config.IMAGE_SHAPE,
config.NUM_CLASSES)
self.config = config
for m in chain(self.classifier.modules(), self.mask.modules()):
if isinstance(m, nn.Conv2d):
nn.init.xavier_uniform(m.weight)
if m.bias is not None:
m.bias.data.zero_().cuda()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.01)
m.bias.data.zero_()
return
def classify(rgb_img):
config = CocoConfig()
model = modellib.MaskRCNN(mode="inference",
model_dir="./models/",
config=config)
model.load_weights(filepath="./models/mask_rcnn_moles_0030.h5",
by_name=True)
class_names = ['Salmonella', 'Yersinia', 'Shigella', 'Negativo']
r = model.detect([rgb_img])[0]
image, caption = visualize.display_instances(rgb_img, r['rois'],
r['masks'], r['class_ids'],
class_names, r['scores'])
K.clear_session()
image = PIL.Image.fromarray(image, 'RGB')
words = caption.split()
filename = str(uuid.uuid4())
filepath = "./static/img/test/{}.jpg".format(filename)
image.save(filepath)
return json.dumps([{
"caption": words[0],
"percent": words[1],
"image_path": filepath
}],
cls=NumpyEncoder)
def run(input_df):
data = json.loads(input_df)
im = load_image(image64=data)
config = CocoConfig()
model = modellib.MaskRCNN(mode="inference", model_dir="./models/", config=config)
model.load_weights(filepath="./models/mask_rcnn_moles_0090.h5", by_name=True)
class_names = ["BG", "malignant", "benign"]
# predict the mask, bounding box and class of the image
r = model.detect([im])[0]
prediction = None
for idx, val in enumerate(class_names):
if idx == r["class_ids"]:
prediction = val
print(val)
else:
continue
return prediction
default=False,
metavar="<True|False>",
help='Shows scaled perturbation (default=False)',
type=bool)
args = parser.parse_args()
print("Model: ", args.model)
print("Dataset: ", args.dataset)
print("Year: ", args.year)
print("Logs: ", args.logs)
print("Target: ", args.target)
print("Show Perturbation: ", args.show_perturbation)
print("Use Mask: ", args.use_mask)
# print("Auto Download: ", args.download)
config = CocoConfig()
config.display()
# Create model
model = MaskRCNN(config=config, model_dir=args.logs)
if config.GPU_COUNT:
model = model.cuda()
# Select weights file to load
model_path = COCO_MODEL_PATH
# Load weights
print("Loading weights ", model_path)
model.load_weights(model_path)
dataset_train = CocoDataset()
dir_path = os.path.dirname(os.path.realpath(__file__))
MODEL_DIR = dir_path + "/models/"
MODEL_PATH = input(
"Insert the path of your trained model [ Like models/moles.../mask_rcnn_moles_0090.h5 ]: "
)
if os.path.isfile(MODEL_PATH) == False:
raise Exception(MODEL_PATH + " Does not exists")
path_data = input(
"Insert the path of Data [ Link /home/../ISIC-Archive-Downloader/Data/ ] : "
)
if not os.path.exists(path_data):
raise Exception(path_data + " Does not exists")
config = CocoConfig()
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
model.load_weights(MODEL_PATH, by_name=True)
class_names = ["BG", "malignant", "benign"]
all_desc_path = glob.glob(path_data + "Descriptions/ISIC_*")
for filename in os.listdir(path_data + "Descriptions/"):
data = json.load(open(path_data + "/Descriptions/" + filename))
img = cv2.imread(path_data + "Images/" + filename + ".jpeg")
img = cv2.resize(img, (128, 128))
## ground truth of the class
print(data["meta"]["clinical"]["benign_malignant"])
def train(
cfg,
data_cfg,
img_size=416,
resume=False,
epochs=273, # 500200 batches at bs 64, dataset length 117263
batch_size=16,
accumulate=1,
multi_scale=False,
freeze_backbone=False,
transfer=False # Transfer learning (train only YOLO layers)
):
init_seeds()
weights = 'weights' + os.sep
latest = weights + 'latest.pt'
best = weights + 'best.pt'
device = torch_utils.select_device()
if multi_scale:
img_size = 608 # initiate with maximum multi_scale size
opt.num_workers = 0 # bug https://github.com/ultralytics/yolov3/issues/174
else:
torch.backends.cudnn.benchmark = True # unsuitable for multiscale
# Configure run
train_path = parse_data_cfg(data_cfg)['train']
# Initialize model
#model = Darknet(cfg, img_size).to(device)
config = CocoConfig()
mask = Mask(256, config.MASK_POOL_SIZE, config.IMAGE_SHAPE,
config.NUM_CLASSES)
# Optimizer
# optimizer = optim.SGD(model.parameters(), lr=hyp['lr0'], momentum=hyp['momentum'], weight_decay=hyp['weight_decay'])
cutoff = -1 # backbone reaches to cutoff layer
start_epoch = 0
best_loss = float('inf')
'''
#nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255)
if resume: # Load previously saved model
if transfer: # Transfer learning
chkpt = torch.load(weights + 'yolov3-spp.pt', map_location=device)
model.load_state_dict({k: v for k, v in chkpt['model'].items() if v.numel() > 1 and v.shape[0] != 255},
strict=False)
for p in model.parameters():
p.requires_grad = True if p.shape[0] == nf else False
else: # resume from latest.pt
chkpt = torch.load(latest, map_location=device) # load checkpoint
model.load_state_dict(chkpt['model'])
start_epoch = chkpt['epoch'] + 1
if chkpt['optimizer'] is not None:
optimizer.load_state_dict(chkpt['optimizer'])
best_loss = chkpt['best_loss']
del chkpt
else: # Initialize model with backbone (optional)
if '-tiny.cfg' in cfg:
cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
else:
cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
# Scheduler (reduce lr at epochs 218, 245, i.e. batches 400k, 450k)
# lf = lambda x: 1 - x / epochs # linear ramp to zero
# lf = lambda x: 10 ** (-2 * x / epochs) # exp ramp to lr0 * 1e-2
# lf = lambda x: 1 - 10 ** (hyp['lrf'] * (1 - x / epochs)) # inv exp ramp to lr0 * 1e-2
# scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lf, last_epoch=start_epoch - 1)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[218, 245],
gamma=0.1,
last_epoch=start_epoch - 1)
# Plot lr schedule
# y = []
# for _ in range(epochs):
# scheduler.step()
# y.append(optimizer.param_groups[0]['lr'])
# plt.plot(y)
# Dataset
#dataset = LoadImagesAndLabels(train_path, img_size=img_size, augment=True)
# Initialize distributed training
if torch.cuda.device_count() > 1:
dist.init_process_group(backend=opt.backend, init_method=opt.dist_url, world_size=opt.world_size, rank=opt.rank)
model = torch.nn.parallel.DistributedDataParallel(model)
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
else:
sampler = None
# Dataloader
#dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=opt.num_workers,
shuffle=False,
pin_memory=True,
collate_fn=dataset.collate_fn,
sampler=sampler)
# Mixed precision training https://github.com/NVIDIA/apex
# install help: https://github.com/NVIDIA/apex/issues/259
mixed_precision = False
if mixed_precision:
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
# Start training
t = time.time()
model.hyp = hyp # attach hyperparameters to model
model_info(model)
nb = len(dataloader)
results = (0, 0, 0, 0, 0) # P, R, mAP, F1, test_loss
n_burnin = min(round(nb / 5 + 1), 1000) # burn-in batches
os.remove('train_batch0.jpg') if os.path.exists('train_batch0.jpg') else None
os.remove('test_batch0.jpg') if os.path.exists('test_batch0.jpg') else None
'''
for epoch in range(start_epoch, epochs):
#model.train()
print(
('\n%8s%12s' + '%10s' * 7) % ('Epoch', 'Batch', 'xy', 'wh', 'conf',
'cls', 'total', 'nTargets', 'time'))
'''
# Update scheduler
scheduler.step()
# Freeze backbone at epoch 0, unfreeze at epoch 1
if freeze_backbone and epoch < 2:
for name, p in model.named_parameters():
if int(name.split('.')[1]) < cutoff: # if layer < 75
p.requires_grad = False if epoch == 0 else True
'''
mloss = torch.zeros(5).to(device) # mean losses
coco_path = "/data/Huaiyu/huaiyu/coco/"
for i in range(1):
#for i, (imgs, targets, gt_mask, _, _) in enumerate(dataloader):
reader_mask = open(
coco_path +
"mask/train2014/COCO_train2014_000000000009.pickle", "rb")
reader_fmap = open(
coco_path +
"feature_maps/train2014/COCO_train2014_000000000009.pickle",
"rb")
pk_mask = pk.load(reader_mask)
pk_fmap = pk.load(reader_fmap)
gt_mask = torch.tensor(pk_mask[0]).to(device)
gt_bbox = torch.tensor(pk_mask[1]).to(device)
pred = torch.tensor(pk_fmap[3][:, :4]).to(device)
feature_maps = pk_fmap[:3]
for idx, item in enumerate(feature_maps):
feature_maps[idx] = torch.tensor(item).to(device)
'''
imgs = imgs.to(device)
targets = targets.to(device)
gt_mask = gt_mask.to(device)
print('imgs:', imgs.shape)
print('targets:', targets.shape)
print('gt_mask:', gt_mask.shape)
nt = len(targets)
# if nt == 0: # if no targets continue
# continue
# Plot images with bounding boxes
if epoch == 0 and i == 0:
plot_images(imgs=imgs, targets=targets, fname='train_batch0.jpg')
# SGD burn-in
if epoch == 0 and i <= n_burnin:
lr = hyp['lr0'] * (i / n_burnin) ** 4
for x in optimizer.param_groups:
x['lr'] = lr
'''
# Run model
#pred, feature_map = model(imgs)
# print('feature map:', len(feature_map))
# print('pred:', len(pred))
print('gt_bbox:', gt_bbox.shape)
print('gt_mask:', gt_mask.shape)
print('pred in training', pred)
print('feature map0 in training', feature_maps[0].shape)
print('feature map1 in training', feature_maps[1].shape)
print('feature map2 in training', feature_maps[2].shape)
'''
image_shape = [416, 416, 3]
pred_temp =[]
for i in range(len(pred)):
temp = pred[i].view(1,-1,85)
pred_temp.append(temp)
pred = torch.cat(pred_temp, 1)
# detections = []
boxes = []
num_boxes = len(pred)
for i in range(num_boxes):
detection = non_max_suppression(pred[i], conf_thres=0.5, nms_thres=0.5)
# detections.append(detection)
x1, y1, x2, y2 = detection
box = [num_boxes, (y1, x1, y2, x2)]
boxes.append(box)
boxes = np.array(boxes)
inputs_roi = [boxes, feature_map]
pooled_regions = pyramid_roi_align(inputs_roi, [14, 14], image_shape)
'''
# mask branch
# boxes = [batch_size, num_boxes, (y1, x1, y2, x2)]
# if x1,y1,x2,y2 is m*1 size
# boxes = torch.cat([y1, x1, y2, x2], dim=1)
boxes = pred
# inputs_roi = [boxes, feature_map]
image_shape = [416, 416, 3]
# pooled_regions = pyramid_roi_align(inputs_roi, [14, 14], image_shape)
mrcnn_mask = mask(feature_maps, boxes)
# print(feature_map[0].shape) # torch.Size([1, 255, 13, 13])
# print(feature_map[1].shape) # torch.Size([1, 255, 26, 26])
# print(feature_map[2].shape) # torch.Size([1, 255, 52, 52])
# Compute loss
mask_loss = compute_mrcnn_mask_loss(target_masks, target_class_ids,
mrcnn_mask)
loss, loss_items = compute_loss(pred, targets, model)
loss = loss + mask_loss
if torch.isnan(loss):
print('WARNING: nan loss detected, ending training')
return results
# Compute gradient
if mixed_precision:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# Accumulate gradient for x batches before optimizing
if (i + 1) % accumulate == 0 or (i + 1) == nb:
optimizer.step()
optimizer.zero_grad()
# Update running mean of tracked metrics
mloss = (mloss * i + loss_items) / (i + 1)
# Print batch results
s = ('%8s%12s' +
'%10.3g' * 7) % ('%g/%g' % (epoch, epochs - 1), '%g/%g' %
(i, nb - 1), *mloss, nt, time.time() - t)
t = time.time()
print(s)
# Multi-Scale training (320 - 608 pixels) every 10 batches
if multi_scale and (i + 1) % 10 == 0:
dataset.img_size = random.choice(range(10, 20)) * 32
print('multi_scale img_size = %g' % dataset.img_size)
# Calculate mAP (always test final epoch, skip first 5 if opt.nosave)
if not (opt.notest or
(opt.nosave and epoch < 5)) or epoch == epochs - 1:
with torch.no_grad():
results = test.test(cfg,
data_cfg,
batch_size=batch_size,
img_size=img_size,
model=model,
conf_thres=0.1)
# Write epoch results
with open('results.txt', 'a') as file:
file.write(s + '%11.3g' * 5 % results +
'\n') # P, R, mAP, F1, test_loss
# Update best loss
test_loss = results[4]
if test_loss < best_loss:
best_loss = test_loss
# Save training results
save = True and not opt.nosave
if save:
# Create checkpoint
chkpt = {
'epoch':
epoch,
'best_loss':
best_loss,
'model':
model.module.state_dict()
if type(model) is nn.parallel.DistributedDataParallel else
model.state_dict(),
'optimizer':
optimizer.state_dict()
}
# Save latest checkpoint
torch.save(chkpt, latest)
# Save best checkpoint
if best_loss == test_loss:
torch.save(chkpt, best)
# Save backup every 10 epochs (optional)
if epoch > 0 and epoch % 10 == 0:
torch.save(chkpt, weights + 'backup%g.pt' % epoch)
# Delete checkpoint
del chkpt
return results