def post(self):
ret = []
nms_threshold = 0.5
cls_threshold = 0.3
classes = ['Background', 'Knife', 'Horse', 'Human']
parser = reqparse.RequestParser()
parser.add_argument('img', required=True)
args = parser.parse_args()
imgurl = args['img']
response = requests.get(imgurl, stream=True)
img = Image.open(BytesIO(response.content)).convert("RGB")
width, height = img.size
model = SSD()
checkpoint = torch.load(
'C:\\Users\\Elina\\Desktop\\IoU-balanced-Loss-SSD\\trained_models\\SSD.pth'
)
model.load_state_dict(checkpoint["model_state_dict"])
if torch.cuda.is_available():
model.cuda()
model.eval()
dboxes = generate_dboxes()
transformer = SimpleTransformer(dboxes, eval=True)
img, _, _ = transformer(img, torch.zeros(4), torch.zeros(1))
encoder = Encoder(dboxes)
if torch.cuda.is_available():
img = img.cuda()
with torch.no_grad():
ploc, plabel = model(img.unsqueeze(dim=0))
result = encoder.decode_batch(ploc, plabel, nms_threshold, 20)[0]
loc, label, prob = [r.cpu().numpy() for r in result]
best = np.argwhere(prob > cls_threshold).squeeze(axis=1)
loc = loc[best]
label = label[best]
prob = prob[best]
if len(loc) > 0:
loc[:, 0::2] *= width
loc[:, 1::2] *= height
loc = loc.astype(np.int32)
for box, lb, pr in zip(loc, label, prob):
category = classes[lb]
xmin, ymin, xmax, ymax = box
xmin, ymin, xmax, ymax = self.ensure_legal(
xmin, ymin, xmax, ymax, width, height)
ret.append({
'label':
category,
'bbox':
[float(xmin),
float(ymin),
float(xmax),
float(ymax)]
})
return json.dumps(ret), 200
def test_one(path):
model = SSD()
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint["model_state_dict"])
if torch.cuda.is_available():
model.cuda()
model.eval()
dboxes = generate_dboxes()
transformer = SimpleTransformer(dboxes, eval=True)
img = Image.open(path).convert("RGB")
img, _, _ = transformer(img, torch.zeros(4), torch.zeros(1))
encoder = Encoder(dboxes)
if torch.cuda.is_available():
img = img.cuda()
with torch.no_grad():
ploc, plabel = model(img.unsqueeze(dim=0))
result = encoder.decode_batch(ploc, plabel, nms_threshold, 20)[0]
loc, label, prob = [r.cpu().numpy() for r in result]
best = np.argwhere(prob > cls_threshold).squeeze(axis=1)
loc = loc[best]
label = label[best]
prob = prob[best]
output_img = cv2.imread(path)
if len(loc) > 0:
height, width, _ = output_img.shape
loc[:, 0::2] *= width
loc[:, 1::2] *= height
loc = loc.astype(np.int32)
for box, lb, pr in zip(loc, label, prob):
category = classes[lb]
color = colors[lb]
xmin, ymin, xmax, ymax = box
xmin, ymin, xmax, ymax = ensure_legal(xmin, ymin, xmax, ymax,
width, height)
cv2.rectangle(output_img, (xmin, ymin), (xmax, ymax), color, 2)
text_size = cv2.getTextSize(category + " : %.2f" % pr,
cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
cv2.rectangle(
output_img, (xmin, ymin),
(xmin + text_size[0] + 3, ymin + text_size[1] + 4), color,
-1)
cv2.putText(output_img, category + " : %.2f" % pr,
(xmin, ymin + text_size[1] + 4),
cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)
output = output_path + "{}_prediction.jpg".format(
os.path.splitext(os.path.split(path)[1])[0])
cv2.imwrite(output, output_img)
def main(opt):
if torch.cuda.is_available():
print('Will compute using CUDA')
# torch.distributed.init_process_group(backend='nccl', init_method='env://')
# num_gpus = torch.distributed.get_world_size()
num_gpus = 1
torch.cuda.manual_seed(123)
else:
torch.manual_seed(123)
num_gpus = 1
train_params = {
"batch_size": opt.batch_size * num_gpus,
"shuffle": True,
"drop_last": False,
"num_workers": opt.num_workers,
"collate_fn": collate_fn
}
test_params = {
"batch_size": opt.batch_size * num_gpus,
"shuffle": True,
"drop_last": False,
"num_workers": opt.num_workers,
"collate_fn": collate_fn
}
dboxes = generate_dboxes()
model = SSD()
train_set = OIDataset(SimpleTransformer(dboxes))
train_loader = DataLoader(train_set, **train_params)
test_set = OIDataset(SimpleTransformer(dboxes, eval=True), train=False)
test_loader = DataLoader(test_set, **test_params)
encoder = Encoder(dboxes)
opt.lr = opt.lr * num_gpus * (opt.batch_size / 32)
criterion = Loss(dboxes)
optimizer = torch.optim.SGD(model.parameters(),
lr=opt.lr,
momentum=opt.momentum,
weight_decay=opt.weight_decay,
nesterov=True)
scheduler = MultiStepLR(optimizer=optimizer,
milestones=opt.multistep,
gamma=0.1)
if torch.cuda.is_available():
model.cuda()
criterion.cuda()
model = torch.nn.DataParallel(model)
if os.path.isdir(opt.log_path):
shutil.rmtree(opt.log_path)
os.makedirs(opt.log_path)
if not os.path.isdir(opt.save_folder):
os.makedirs(opt.save_folder)
checkpoint_path = os.path.join(opt.save_folder, "SSD.pth")
writer = SummaryWriter(opt.log_path)
if os.path.isfile(checkpoint_path):
checkpoint = torch.load(checkpoint_path)
first_epoch = checkpoint["epoch"] + 1
model.module.load_state_dict(checkpoint["model_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler"])
optimizer.load_state_dict(checkpoint["optimizer"])
# evaluate(model, test_loader, encoder, opt.nms_threshold)
else:
first_epoch = 0
for epoch in range(first_epoch, opt.epochs):
train(model, train_loader, epoch, writer, criterion, optimizer,
scheduler)
evaluate(model, test_loader, encoder, opt.nms_threshold)
checkpoint = {
"epoch": epoch,
"model_state_dict": model.module.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict()
}
torch.save(checkpoint, checkpoint_path)