def run(self, image_or_path_or_tensor, meta=None):
if isinstance(image_or_path_or_tensor, np.ndarray):
image = image_or_path_or_tensor
elif type(image_or_path_or_tensor) == type(''):
image = cv2.imread(image_or_path_or_tensor)
detections = []
cost = 0
for scale in self.scales:
images, meta = self.pre_process(image, scale, meta)
images = images.to(device)
# print('input shape: {}'.format(images.shape))
torch.cuda.synchronize()
tic = time.time()
_, dets = self.process(images, return_time=True)
cost = time.time() - tic
print('cost: {}, fps: {}'.format(cost, 1 / cost))
torch.cuda.synchronize()
dets = self.post_process(dets, meta, scale)
torch.cuda.synchronize()
detections.append(dets)
print(detections)
res = visualize_det_cv2(image, detections[0], coco_label_map_list[1:],
0.3)
cv2.putText(res, 'fps: {0:.4f}'.format(1 / cost), (30, 30),
cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 0, 255), 2)
return res
def demo(v_f):
cfg = Config.fromfile(config_f)
anchor_config = anchors(cfg)
priorbox = PriorBox(anchor_config)
net = build_net('test',
size=cfg.model.input_size,
config=cfg.model.m2det_config)
init_net(net, cfg, checkpoint_path)
net.eval().to(device)
with torch.no_grad():
priors = priorbox.forward().to(device)
_preprocess = BaseTransform(
cfg.model.input_size, cfg.model.rgb_means, (2, 0, 1))
detector = Detect(cfg.model.m2det_config.num_classes,
cfg.loss.bkg_label, anchor_config)
logging.info('detector initiated.')
cap = cv2.VideoCapture(v_f)
logging.info('detect on: {}'.format(v_f))
logging.info('video width: {}, height: {}'.format(int(cap.get(3)), int(cap.get(4))))
out_video = cv2.VideoWriter("result.mp4", cv2.VideoWriter_fourcc(*'MJPG'), 24, (int(cap.get(3)), int(cap.get(4))))
while True:
ret, image = cap.read()
if not ret:
out_video.release()
cv2.destroyAllWindows()
cap.release()
break
w, h = image.shape[1], image.shape[0]
img = _preprocess(image).unsqueeze(0).to(device)
scale = torch.Tensor([w, h, w, h])
out = net(img)
boxes, scores = detector.forward(out, priors)
boxes = (boxes[0]*scale).cpu().numpy()
scores = scores[0].cpu().numpy()
allboxes = []
for j in range(1, cfg.model.m2det_config.num_classes):
inds = np.where(scores[:, j] > cfg.test_cfg.score_threshold)[0]
if len(inds) == 0:
continue
c_bboxes = boxes[inds]
c_scores = scores[inds, j]
c_dets = np.hstack((c_bboxes, c_scores[:, np.newaxis])).astype(
np.float32, copy=False)
soft_nms = cfg.test_cfg.soft_nms
# min_thresh, device_id=0 if cfg.test_cfg.cuda else None)
keep = nms(c_dets, cfg.test_cfg.iou, force_cpu=soft_nms)
keep = keep[:cfg.test_cfg.keep_per_class]
c_dets = c_dets[keep, :]
allboxes.extend([_.tolist()+[j] for _ in c_dets])
if len(allboxes) > 0:
allboxes = np.array(allboxes)
# [boxes, scores, label_id] -> [id, score, boxes] 0, 1, 2, 3, 4, 5
allboxes = allboxes[:, [5, 4, 0, 1, 2, 3]]
logging.info('allboxes shape: {}'.format(allboxes.shape))
res = visualize_det_cv2(image, allboxes, classes=classes, thresh=0.2)
# res = visualize_det_cv2_fancy(image, allboxes, classes=classes, thresh=0.2, r=4, d=6)
cv2.imshow('rr', res)
out_video.write(res)
cv2.waitKey(1)