class deepsort_rbc():
def __init__(self):
self.metric = nn_matching.NearestNeighborDistanceMetric(
"cka", .5, 30) #metric, matching_threshold, budget
self.tracker = Tracker(self.metric)
def reset_tracker(self):
self.tracker = Tracker(self.metric)
#Deep sort needs the format `top_left_x, top_left_y, width,height
def format_boxes(self, out_boxes):
for b in range(len(out_boxes)):
out_boxes[b][2] = out_boxes[b][2] - out_boxes[b][0]
out_boxes[b][3] = out_boxes[b][3] - out_boxes[b][1]
return out_boxes
def run_deep_sort(self, frame, out_boxes, out_scores, classIDs, scales,
ids, features):
if out_boxes == []:
self.tracker.predict()
print('No detections')
trackers = self.tracker.tracks
return trackers
wh = np.flip(frame.shape[0:2])
out_boxes[:, 0:2] = (out_boxes[:, 0:2] * wh).astype(float)
out_boxes[:, 2:4] = (out_boxes[:, 2:4] * wh).astype(float)
# Give format boxes
out_boxes = self.format_boxes(out_boxes)
# Create a detection object to parse in deepsort
dets = [Detection(bbox, score, feature) \
for bbox, score, feature in\
zip(out_boxes, out_scores, features)]
outboxes = np.array([d.tlwh for d in dets])
outscores = np.array([d.confidence for d in dets])
# Non max suppression including confidence
# works best using pixel coordinates, 0.3<=overlap<=0.5
indices = non_max_suppression(outboxes, 0.4, outscores)
dets = [dets[i] for i in indices]
classIDs = [classIDs[i] for i in indices]
scales = [scales[i] for i in indices]
ids = [ids[i] for i in indices]
detections_class = (dets, classIDs, scales, ids)
# DeepSort cycle
self.tracker.predict()
self.tracker.update(dets)
return self.tracker, detections_class
class deepsort_rbc():
def __init__(self):
#loading this encoder is slow, should be done only once.
#self.encoder = generate_detections.create_box_encoder("deep_sort/resources/networks/mars-small128.ckpt-68577")
self.encoder = vehicle_encoder('weights/vehicle_vgg.h5')
self.metric = nn_matching.NearestNeighborDistanceMetric(
"cosine", .5, 100)
self.tracker = Tracker(self.metric)
def reset_tracker(self):
self.tracker = Tracker(self.metric)
#Older yolo used to give y,x,h,w.
"""
def format_yolo_output( self,out_boxes):
out_boxes=np.array([out_boxes[:,1],out_boxes[:,0],\
out_boxes[:,3]-out_boxes[:,1],out_boxes[:,2]-out_boxes[:,0]])
out_boxes=out_boxes.T
return out_boxes
"""
#Current Yolo gives x_center,y_center,w,h
#Deep sort needs the format `top_left_x, top_left_y, width,height
def format_yolo_output(self, out_boxes):
for b in range(len(out_boxes)):
out_boxes[b][0] = out_boxes[b][0] - out_boxes[b][2] / 2
out_boxes[b][1] = out_boxes[b][1] - out_boxes[b][3] / 2
return out_boxes
def run_deep_sort(self, frame, out_scores, out_boxes, out_classes):
out_boxes = self.format_yolo_output(out_boxes)
if out_boxes == []:
self.tracker.predict()
trackers = self.tracker.tracks
return trackers
detections = np.array(out_boxes)
#features = self.encoder(frame, detections.copy())
features = self.encoder.extract_features(frame, detections)
#print(frame.shape)
detections = [Detection(bbox, score, feature,classname) \
for bbox,score, feature,classname in\
zip(detections,out_scores, features,out_classes)]
outboxes = np.array([d.tlwh for d in detections])
outscores = np.array([d.confidence for d in detections])
indices = prep.non_max_suppression(outboxes, 0.8, outscores)
detections = [detections[i] for i in indices]
self.tracker.predict()
self.tracker.update(detections)
trackers = self.tracker.tracks
return trackers
class track_deepsort(object):
def __init__(self):
self.metric = nn_matching.NearestNeighborDistanceMetric(
"cosine", 0.2, 100)
self.tracker = Tracker(self.metric)
self.encoder = generate_detections.create_box_encoder(
os.path.abspath(
"deep_sort/resources/networks/mars-small128.ckpt-68577"))
def return_tracking_id(self, bboxes, frame):
detections = []
scores = []
new_bboxes = []
nms_max_overlap = 0.1
if type(frame) is not np.ndarray:
imgcv = cv2.imread(frame)
else:
imgcv = frame
h, w, _ = frame.shape
thick = int((h + w) // 100)
for b in bboxes[0]:
left, top, right, bot, confidence, _, _ = b
detections.append(
np.array([left, top, right - left,
bot - top]).astype(np.float64))
scores.append(confidence)
detections = np.array(detections)
if detections.shape[0] == 0:
return frame
scores = np.array(scores)
features = self.encoder(imgcv, detections.copy())
detections = [
Detection(bbox, score, feature)
for bbox, score, feature in zip(detections, scores, features)
] # Run non-maxima suppression.
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
indices = prep.non_max_suppression(boxes, nms_max_overlap, scores)
detections = [detections[i] for i in indices]
self.tracker.predict()
self.tracker.update(detections)
trackers = self.tracker.tracks
for track in trackers:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
id_num = str(track.track_id)
new_bboxes.append(
(id_num, int(bbox[0]), int(bbox[1]), (int(bbox[2])),
int(bbox[3]), track.confidence))
print(new_bboxes)
return new_bboxes
class deepsort_rbc():
def __init__(
self,
wt_path="/gdrive/MyDrive/car_tracking/car_tracking/object_tracking/ckpts/model640.pt"
):
#loading this encoder is slow, should be done only once.
#self.encoder = generate_detections.create_box_encoder("deep_sort/resources/networks/mars-small128.ckpt-68577")
self.encoder = torch.load(wt_path)
self.encoder = self.encoder.cuda()
self.encoder = self.encoder.eval()
print("Deep sort model loaded from path: ", wt_path)
self.metric = nn_matching.NearestNeighborDistanceMetric(
"cosine", .5, 100)
self.tracker = Tracker(self.metric)
self.gaussian_mask = get_gaussian_mask().cuda()
self.transforms = torchvision.transforms.Compose([ \
torchvision.transforms.ToPILImage(),\
torchvision.transforms.Resize((128,128)),\
torchvision.transforms.ToTensor()])
def reset_tracker(self):
self.tracker = Tracker(self.metric)
#Deep sort needs the format `top_left_x, top_left_y, width,height
def format_yolo_output(self, out_boxes):
for b in range(len(out_boxes)):
out_boxes[b][0] = out_boxes[b][0] - out_boxes[b][2] / 2
out_boxes[b][1] = out_boxes[b][1] - out_boxes[b][3] / 2
return out_boxes
def pre_process(self, frame, detections):
transforms = torchvision.transforms.Compose([ \
torchvision.transforms.ToPILImage(),\
torchvision.transforms.Resize((128,128)),\
torchvision.transforms.ToTensor()])
crops = []
for d in detections:
for i in range(len(d)):
if d[i] < 0:
d[i] = 0
img_h, img_w, img_ch = frame.shape
xmin, ymin, w, h = d
if xmin > img_w:
xmin = img_w
if ymin > img_h:
ymin = img_h
xmax = xmin + w
ymax = ymin + h
ymin = abs(int(ymin))
ymax = abs(int(ymax))
xmin = abs(int(xmin))
xmax = abs(int(xmax))
try:
crop = frame[ymin:ymax, xmin:xmax, :]
crop = transforms(crop)
crops.append(crop)
except:
continue
crops = torch.stack(crops)
return crops
def extract_features_only(self, frame, coords):
for i in range(len(coords)):
if coords[i] < 0:
coords[i] = 0
img_h, img_w, img_ch = frame.shape
xmin, ymin, w, h = coords
if xmin > img_w:
xmin = img_w
if ymin > img_h:
ymin = img_h
xmax = xmin + w
ymax = ymin + h
ymin = abs(int(ymin))
ymax = abs(int(ymax))
xmin = abs(int(xmin))
xmax = abs(int(xmax))
crop = frame[ymin:ymax, xmin:xmax, :]
#crop = crop.astype(np.uint8)
#print(crop.shape,[xmin,ymin,xmax,ymax],frame.shape)
crop = self.transforms(crop)
crop = crop.cuda()
gaussian_mask = self.gaussian_mask
input_ = crop * gaussian_mask
input_ = torch.unsqueeze(input_, 0)
features = self.encoder.forward_once(input_)
features = features.detach().cpu().numpy()
corrected_crop = [xmin, ymin, xmax, ymax]
return features, corrected_crop
def run_deep_sort(self, frame, out_scores, out_boxes):
if out_boxes == []:
self.tracker.predict()
print('No detections')
trackers = self.tracker.tracks
return trackers
detections = np.array(out_boxes)
#features = self.encoder(frame, detections.copy())
processed_crops = self.pre_process(frame, detections).cuda()
processed_crops = self.gaussian_mask * processed_crops
features = self.encoder.forward_once(processed_crops)
features = features.detach().cpu().numpy()
if len(features.shape) == 1:
features = np.expand_dims(features, 0)
dets = [Detection(bbox, score, feature) \
for bbox,score, feature in\
zip(detections,out_scores, features)]
outboxes = np.array([d.tlwh for d in dets])
outscores = np.array([d.confidence for d in dets])
indices = prep.non_max_suppression(outboxes, 0.8, outscores)
dets = [dets[i] for i in indices]
self.tracker.predict()
self.tracker.update(dets)
return self.tracker, dets
class DeepSortNode(Observer, Subject):
objectBoundingBoxes = []
objectIds = []
image = []
END = False
def __init__(self, encoderPath, applyMask=False):
self.metric = nn_matching.NearestNeighborDistanceMetric(
"cosine", 0.9, 100)
self.tracker = Tracker(self.metric,
max_iou_distance=0.9,
max_age=50,
n_init=3,
_lambda=0.3)
self.encoder = generate_detections.create_box_encoder(
encoderPath, applyMask=applyMask)
self.applyMask = applyMask
def update(self, subject):
if (subject.END):
self.END = True
self.notify()
return
# update tracker with detection from detector
boundBoxes = subject.rois
confidences = subject.scores
if (self.applyMask):
features = self.encoder(subject.image, np.array(boundBoxes),
subject.masks)
else:
features = self.encoder(subject.image, np.array(boundBoxes))
detections = [
Detection(bbox, confidence,
feature) for bbox, confidence, feature in zip(
boundBoxes, confidences, features)
]
self.tracker.predict()
self.tracker.update(detections)
# extract bounding boxes and Ids
self.image = subject.image
self.objectBoundingBoxes = []
self.objectIds = []
tracks = self.tracker.tracks
for track in tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
self.objectBoundingBoxes.append(track.to_tlbr())
self.objectIds.append(str(track.track_id))
# notify deep sort event listeners
self.notify()
def notify(self):
print("DeepSortNode: update people tracking")
for observer in self.observers:
observer.update(self)
class Deepsort_original(object):
def __init__(self, wt_path='ckpts/cos_metric_net.pt', wt_type='cosine'):
if wt_type == 'cosine':
exclude_keys = ['weights',
'scale'] # these params useless in eval stage
self.encoder = CosineMetricNet(num_classes=-1,
add_logits=False).cuda()
try:
ckpt: collections.OrderedDict = torch.load(wt_path)
# type: ckpt['model_state_dict']: dict
ckpt['model_state_dict'] = {
k: v
for k, v in ckpt['model_state_dict'].items()
if k not in exclude_keys
}
self.encoder.load_state_dict(ckpt['model_state_dict'],
strict=False)
except KeyError as e:
s = "Model loaded(%s) is not compatible with the definition, please check!" % wt_path
raise KeyError(s) from e
elif wt_type == 'siamese':
self.encoder = torch.load(wt_path)
else:
raise NotImplementedError
self.encoder = self.encoder.eval()
print("Deep sort model loaded")
self.metric = nn_matching.NearestNeighborDistanceMetric(
"cosine", .5, 100)
self.tracker = Tracker(self.metric)
self.transforms = torchvision.transforms.Compose([
torchvision.transforms.ToPILImage(),
torchvision.transforms.Resize((128, 64)),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
def reset_tracker(self):
self.tracker = Tracker(self.metric)
# Deep sort needs the format `top_left_x, top_left_y, width,height
def format_yolo_output(self, out_boxes):
for b in range(len(out_boxes)):
out_boxes[b][0] = out_boxes[b][0] - out_boxes[b][2] / 2
out_boxes[b][1] = out_boxes[b][1] - out_boxes[b][3] / 2
return out_boxes
def pre_process(self, frame, detections):
transforms = torchvision.transforms.Compose([
torchvision.transforms.ToPILImage(),
torchvision.transforms.Resize((128, 64)),
torchvision.transforms.ToTensor(),
])
crops = []
for d in detections:
for i in range(len(d)):
if d[i] < 0:
d[i] = 0
img_h, img_w, img_ch = frame.shape
xmin, ymin, w, h = d
if xmin > img_w:
xmin = img_w
if ymin > img_h:
ymin = img_h
xmax = xmin + w
ymax = ymin + h
ymin = abs(int(ymin))
ymax = abs(int(ymax))
xmin = abs(int(xmin))
xmax = abs(int(xmax))
try:
crop = frame[ymin:ymax, xmin:xmax, :]
crop = transforms(crop)
crops.append(crop)
except:
continue
crops = torch.stack(crops)
return crops
def extract_features_only(self, frame, coords):
for i in range(len(coords)):
if coords[i] < 0:
coords[i] = 0
img_h, img_w, img_ch = frame.shape
xmin, ymin, w, h = coords
if xmin > img_w:
xmin = img_w
if ymin > img_h:
ymin = img_h
xmax = xmin + w
ymax = ymin + h
ymin = abs(int(ymin))
ymax = abs(int(ymax))
xmin = abs(int(xmin))
xmax = abs(int(xmax))
crop = frame[ymin:ymax, xmin:xmax, :]
# crop = crop.astype(np.uint8)
# print(crop.shape,[xmin,ymin,xmax,ymax],frame.shape)
crop = self.transforms(crop)
crop = crop.cuda()
gaussian_mask = self.gaussian_mask
input_ = crop * gaussian_mask
input_ = torch.unsqueeze(input_, 0)
features = self.encoder.forward_once(input_)
features = features.detach().cpu().numpy()
corrected_crop = [xmin, ymin, xmax, ymax]
return features, corrected_crop
def run_deep_sort(self, frame, out_scores, out_boxes):
if out_boxes == []:
self.tracker.predict()
print('No detections')
trackers = self.tracker.tracks
return trackers
detections = np.array(out_boxes)
# features = self.encoder(frame, detections.copy())
processed_crops = self.pre_process(frame, detections).cuda()
# processed_crops = self.gaussian_mask * processed_crops
features = self.encoder.forward_once(processed_crops)
features = features.detach().cpu().numpy()
if len(features.shape) == 1:
features = np.expand_dims(features, 0)
dets = [Detection(bbox, score, feature) \
for bbox, score, feature in \
zip(detections, out_scores, features)]
self.tracker.predict()
self.tracker.update(dets)
return self.tracker, dets
class yact_node:
def __init__(self, debug=0):
self.debug = debug
self.frameCount = 0
self.timePrev = time.time()
#region DEEPSORT
metric = nn_matching.NearestNeighborDistanceMetric(
'cosine', matching_threshold=0.2, budget=100)
self.tracker = Tracker(metric)
absFilePath = os.path.abspath(__file__)
fileDir = os.path.dirname(absFilePath)
filePathModel = os.path.join(
fileDir, 'deep_sort/resources/networks/mars-small128.pb')
self.encoder = gdet.create_box_encoder(filePathModel, batch_size=1)
#endregion
self.detectionAndId = []
self.subscriberImageDetections = rospy.Subscriber(
'yolo_detector/output/compresseddetections',
CompressedImageAndBoundingBoxes,
self.callback,
queue_size=1)
self.publisherDetectionID = rospy.Publisher('yact/output/detectionids',
DetectionAndID,
queue_size=1)
def callback(self, msg):
self.img = cv2.imdecode(np.fromstring(msg.data, np.uint8), 1)
lstDetections = msg.bounding_boxes
lstDetsDeepSort = []
if (lstDetections):
for det in lstDetections:
if det.Class == 'person':
# Deep Sort Bounding Boxes
# Use the format TOP LEFT WIDTH HEIGHT (tlwh)
dsWidth = det.xmax - det.xmin
dsHeight = det.ymax - det.ymin
lstDetsDeepSort.append(
[det.xmin, det.ymin, dsWidth, dsHeight])
self.detectionAndId = []
#region DEEPSORT
if self.frameCount % 3 == 0:
features = self.encoder(self.img, lstDetsDeepSort)
# Create DeepSort detections
trackedObjects = [
Detection(bbox, 1.0, feature)
for bbox, feature in zip(lstDetsDeepSort, features)
]
self.tracker.predict()
self.tracker.update(trackedObjects)
#endregion
self.displayDetections()
msgDetectionAndID = DetectionAndID()
msgDetectionAndID.header = msg.header
msgDetectionAndID.detections = self.detectionAndId
#region DISPLAY
if self.frameCount % 10 == 0:
timer = time.time() - self.timePrev
self.timePrev = time.time()
self.intFPS = int(10 / timer)
# Print FPS
cv2.putText(self.img, 'FPS: {}'.format(self.intFPS),
(self.img.shape[1] - 100, 25), cv2.FONT_HERSHEY_SIMPLEX,
0.5, (255, 0, 0), 2)
cv2.imshow('yact: people_tracker.py', self.img)
cv2.waitKey(3)
self.frameCount += 1
#endregion
#region DEBUG
def displayDetections(self):
for track in self.tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
# Construct message
msgBbid = BoundingBoxID()
msgBbid.boundingBox.xmin = int(bbox[0])
msgBbid.boundingBox.ymin = int(bbox[1])
msgBbid.boundingBox.xmax = int(bbox[2])
msgBbid.boundingBox.ymax = int(bbox[3])
msgBbid.id = int(track.track_id)
self.detectionAndId.append(msgBbid)
# Draw bounding box and label
cv2.rectangle(self.img, (int(bbox[0]), int(bbox[1])),
(int(bbox[2]), int(bbox[3])), (255, 0, 0), 2)
labelText = 'ID: {}'.format(track.track_id)
cv2.putText(self.img, labelText, (int(bbox[0]), int(bbox[1]) - 3),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
def efficientDet_video_inference(video_src,compound_coef = 0,force_input_size=None,
frame_skipping = 3,
threshold=0.2,out_path=None,imshow=False,
display_fps=False):
#deep-sort variables
# Definition of the parameters
max_cosine_distance = 0.3
nn_budget = None
nms_max_overlap = 1.0
model_filename = '/home/shaheryar/Desktop/Projects/Football-Monitoring/deep_sort/model_weights/mars-small128.pb'
encoder = gdet.create_box_encoder(model_filename, batch_size=1)
metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
tracker = Tracker(metric,n_init=5)
# efficientDet-pytorch variables
iou_threshold = 0.4
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size
# load model
model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list))
model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
# Video capture
cap = cv2.VideoCapture(video_src)
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
fourcc = cv2.VideoWriter_fourcc(*'MPEG')
fps = cap.get(cv2.CAP_PROP_FPS)
print("Video fps",fps)
if(out_path is not None):
outp = cv2.VideoWriter(out_path, fourcc, fps, (frame_width, frame_height))
i=0
start= time.time()
current_frame_fps=0
while True:
ret, frame = cap.read()
if not ret:
break
t1=time.time()
if (frame_skipping==0 or i%frame_skipping==0):
# if(True):
# frame preprocessing (running detections)
ori_imgs, framed_imgs, framed_metas, t1 = preprocess_video(frame, width=input_size, height=input_size)
if use_cuda:
x = torch.stack([fi.cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
# model predict
t1=time.time()
with torch.no_grad():
features, regression, classification, anchors = model(x)
out = postprocess(x,
anchors, regression, classification,
regressBoxes, clipBoxes,
threshold, iou_threshold)
# Post processing
out = invert_affine(framed_metas, out)
# decoding bbox ,object name and scores
boxes,classes,scores =decode_predictions(out[0])
org_boxes = boxes.copy()
t2 = time.time() - t1
# feature extraction for deep sort
boxes = [convert_bbox_to_deep_sort_format(frame.shape, b) for b in boxes]
features = encoder(frame,boxes)
detections = [Detection(bbox, 1.0, feature) for bbox, feature in zip(boxes, features)]
boxes = np.array([d.tlwh for d in detections])
# print(boxes)
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(boxes, nms_max_overlap, scores)
detections = [detections[i] for i in indices]
tracker.predict()
tracker.update(detections)
i = i + 1
img_show=frame.copy()
for j in range(len(org_boxes)):
img_show =drawBoxes(img_show,org_boxes[j],(255,255,0),str(tracker.tracks[j].track_id))
for track in tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
x1=int(bbox[0])
y1 = int(bbox[1])
x2 = int(bbox[2])
y2=int(bbox[3])
roi= frame[y1:y2,x1:x2]
cv2.rectangle(img_show, (x1, y1), (x2, y2), update_color_association(roi, track.track_id), 2)
cv2.putText(img_show, str(track.track_id), (x1, y1), 0, 5e-3 * 100, (255, 255, 0), 1)
if display_fps:
current_frame_fps=1/t2
else:
current_frame_fps=0
cv2.putText(img_show, 'FPS: {0:.2f}'.format(current_frame_fps), (30, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
(255, 255, 0),
2, cv2.LINE_AA)
if (i % int(fps) == 0):
print("Processed ", str(int(i / fps)), "seconds")
print("Time taken",time.time()-start)
# print(color_dict)
if imshow:
img_show=cv2.resize(img_show,(0,0),fx=0.75,fy=0.75)
cv2.imshow('Frame',img_show)
# Press Q on keyboard to exit
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if out_path is not None:
outp.write(img_show)
cap.release()
outp.release()
# put rectangle
cv2.rectangle(flipped,
tuple(pred_bbox_pixel[i, :2].astype(int)),
tuple(pred_bbox_pixel[i, 2:].astype(int)),
(255, 0, 0), 2)
for ldx, ldy, color in zip(
pred_ldmk_pixel[i][0::2].astype(int),
pred_ldmk_pixel[i][1::2].astype(int), [(255, 0, 0),
(0, 255, 0),
(0, 0, 255),
(255, 255, 0),
(255, 0, 255)]):
cv2.circle(flipped, (ldx, ldy), 3, color, 2)
ds_tracker.predict()
ds_tracker.update(detections)
for track in ds_tracker.tracks:
cv2.rectangle(flipped, tuple(track.to_tlbr().astype(int)[:2]),
tuple(track.to_tlbr().astype(int)[2:]),
(0, 0, 255), 2)
# print(track.to_tlbr())
cv2.imshow('detection', flipped)
# print fps
print("fps:", 1 / (time.time() - start))
if cv2.waitKey(1) & 0xFF == ord('q'):
break
_, frame = video_capturer.read()
bboxes = personDetector.Detect(frame)
# Draw detections
for bbox in bboxes:
frame = personDetector.DrawBox(
frame, bbox, videoSize, thickness=5)
if deepSort:
bboxes = personDetector.returnJustPersonBox(bboxes)
features = encoder(frame, np.array(
personDetector.bboxToWH(bboxes)))
detections = [Detection(bbox, 1.0, feature) for bbox, feature in zip(
personDetector.bboxToWH(bboxes), features)]
# Call tracker
tracker.predict()
tracker.update(detections)
for track in tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
frame = personDetector.DrawBox(frame, personDetector.bboxToXY(
bbox, videoSize), videoSize, color=(255, 255, 255), thickness=2)
cv2.putText(frame, str(track.track_id), (int(
bbox[0]), int(bbox[1])), 0, 5e-3 * 200, (0, 255, 0), 2)
elif sort:
bboxes = personDetector.returnJustPersonBox(bboxes)
track_bbs_ids = mot_tracker.update(np.array(bboxes))
# Draw tracker
def vis_detections_video(im,
class_name,
dets,
csv_file,
csv,
frame_id,
thresh=0.5):
"""Draw detected bounding boxes."""
nms_max_overlap = 0.6
metric = nn_matching.NearestNeighborDistanceMetric("cosine", 0.2, 100)
tracker = Tracker(metric)
detections = []
scores = []
h, w, _ = im.shape
thick = int((h + w) // 300)
inds = np.where(dets[:, -1] >= thresh)[0]
if len(inds) == 0:
return im
for i in inds:
scores.append(dets[i, -1])
for i in inds:
bbox = dets[i, :4]
boxResults = process_box(bbox, scores, h, w, thresh)
if boxResults is None:
continue
left, right, top, bot, mess, max_indx, confidence = boxResults
detections.append(
np.array([left, top, right - left, bot - top]).astype(np.float64))
scores.append(confidence)
scores = np.array(scores)
detections = np.array(detections)
features = deep_sort_encode(im, detections.copy())
detections = [
Detection(bbox, score, feature)
for bbox, score, feature in zip(detections, scores, features)
]
# Run non-maxima suppression.
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
indices = prep.non_max_suppression(boxes, nms_max_overlap, scores)
detections = [detections[i] for i in indices]
tracker.predict()
tracker.update(detections)
trackers = tracker.tracks
for track in trackers:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
id_num = str(track.track_id)
csv.writerow([
frame_id, id_num,
int(bbox[0]),
int(bbox[1]),
int(bbox[2]) - int(bbox[0]),
int(bbox[3]) - int(bbox[1])
])
csv_file.flush()
cv2.rectangle(im, (int(bbox[0]), int(bbox[1])),
(int(bbox[2]), int(bbox[3])), (0, 255, 255), thick // 3)
cv2.putText(im, id_num, (int(bbox[0]), int(bbox[1]) - 12), 0, 1e-3 * h,
(255, 255, 255), thick // 6)
# cv2.rectangle(im,(bbox[0],bbox[1]),(bbox[2],bbox[3]),(0,0,255),2)
# cv2.rectangle(im,(int(bbox[0]),int(bbox[1])-10),(int(bbox[0]+200),int(bbox[1])+10),(10,10,10),-1)
# cv2.putText(im, id_num,(int(bbox[0]),int(bbox[1]-2)),cv2.FONT_HERSHEY_SIMPLEX,.45,(255,255,255))#,cv2.CV_AA)
return im