jinql
/
Yolov5-deepsort-inference
mirror of https://github.com/Sharpiless/Yolov5-deepsort-inference
1
0
Fork 0
Code Issues Releases Wiki Activity

fix zero detection bug

pull/28/head
Sharpiless 2021-01-07 17:03:54 +08:00 committed by GitHub
parent 7aa228ed2d
commit 8166ef6345
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 74 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

148
tracker.py
View File

@ -1,74 +1,74 @@
from deep_sort.utils.parser import get_config from deep_sort.utils.parser import get_config
from deep_sort.deep_sort import DeepSort from deep_sort.deep_sort import DeepSort
import torch import torch
import cv2 import cv2
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
cfg = get_config() cfg = get_config()
cfg.merge_from_file("deep_sort/configs/deep_sort.yaml") cfg.merge_from_file("deep_sort/configs/deep_sort.yaml")
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True) use_cuda=True)
def plot_bboxes(image, bboxes, line_thickness=None): def plot_bboxes(image, bboxes, line_thickness=None):
# Plots one bounding box on image img # Plots one bounding box on image img
tl = line_thickness or round( tl = line_thickness or round(
0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 # line/font thickness 0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 # line/font thickness
for (x1, y1, x2, y2, cls_id, pos_id) in bboxes: for (x1, y1, x2, y2, cls_id, pos_id) in bboxes:
if cls_id in ['smoke', 'phone', 'eat']: if cls_id in ['smoke', 'phone', 'eat']:
color = (0, 0, 255) color = (0, 0, 255)
else: else:
color = (0, 255, 0) color = (0, 255, 0)
if cls_id == 'eat': if cls_id == 'eat':
cls_id = 'eat-drink' cls_id = 'eat-drink'
c1, c2 = (x1, y1), (x2, y2) c1, c2 = (x1, y1), (x2, y2)
cv2.rectangle(image, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) cv2.rectangle(image, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
tf = max(tl - 1, 1) # font thickness tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(cls_id, 0, fontScale=tl / 3, thickness=tf)[0] t_size = cv2.getTextSize(cls_id, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA) # filled cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA) # filled
cv2.putText(image, '{} ID-{}'.format(cls_id, pos_id), (c1[0], c1[1] - 2), 0, tl / 3, cv2.putText(image, '{} ID-{}'.format(cls_id, pos_id), (c1[0], c1[1] - 2), 0, tl / 3,
[225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
return image return image
def update_tracker(target_detector, image): def update_tracker(target_detector, image):
new_faces = [] new_faces = []
_, bboxes = target_detector.detect(image) _, bboxes = target_detector.detect(image)
bbox_xywh = [] bbox_xywh = []
confs = [] confs = []
# Adapt detections to deep sort input format # Adapt detections to deep sort input format
for x1, y1, x2, y2, _, conf in bboxes: for x1, y1, x2, y2, _, conf in bboxes:
obj = [ obj = [
int((x1+x2)/2), int((y1+y2)/2), int((x1+x2)/2), int((y1+y2)/2),
x2-x1, y2-y1 x2-x1, y2-y1
] ]
bbox_xywh.append(obj) bbox_xywh.append(obj)
confs.append(conf) confs.append(conf)
xywhs = torch.Tensor(bbox_xywh) xywhs = torch.Tensor(bbox_xywh)
confss = torch.Tensor(confs) confss = torch.Tensor(confs)
# Pass detections to deepsort # Pass detections to deepsort
outputs = deepsort.update(xywhs, confss, image) outputs = deepsort.update(xywhs, confss, image)
bboxes2draw = [] bboxes2draw = []
face_bboxes = [] face_bboxes = []
for value in list(outputs): for value in list(outputs):
x1,y1,x2,y2,track_id = value x1,y1,x2,y2,track_id = value
bboxes2draw.append( bboxes2draw.append(
(x1, y1, x2, y2, '', track_id) (x1, y1, x2, y2, '', track_id)
) )
image = plot_bboxes(image, bboxes2draw) image = plot_bboxes(image, bboxes2draw)
return image, new_faces, face_bboxes return image, new_faces, face_bboxes