diff --git a/detect.py b/detect.py index e9a5b91..7f890c4 100644 --- a/detect.py +++ b/detect.py @@ -204,7 +204,7 @@ def detect(save_img=False): print(opt) #check_requirements(exclude=('pycocotools', 'thop')) - if opt.download and not os.path.exists(str(opt.weights)): + if opt.download and not os.path.exists("".join(opt.weights)): print('Model weights not found. Attempting to download now...') download('./') diff --git a/detect_and_track.py b/detect_and_track.py index 620fb19..4d696bd 100644 --- a/detect_and_track.py +++ b/detect_and_track.py @@ -25,7 +25,7 @@ #............................... Bounding Boxes Drawing ............................ """Function to Draw Bounding boxes""" -def draw_boxes(img, bbox, identities=None, categories=None, names=None, save_with_object_id=False, path=None,offset=(0, 0)): +def draw_boxes(img, bbox, identities=None, categories=None, names=None, save_with_object_id=False, path=None, offset=(0, 0), is_label=True): for i, box in enumerate(bbox): x1, y1, x2, y2 = [int(i) for i in box] x1 += offset[0] @@ -35,12 +35,14 @@ def draw_boxes(img, bbox, identities=None, categories=None, names=None, save_wit cat = int(categories[i]) if categories is not None else 0 id = int(identities[i]) if identities is not None else 0 data = (int((box[0]+box[2])/2),(int((box[1]+box[3])/2))) - label = str(id) + ":"+ names[cat] - (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1) - cv2.rectangle(img, (x1, y1), (x2, y2), (255,0,20), 2) - cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), (255,144,30), -1) - cv2.putText(img, label, (x1, y1 - 5),cv2.FONT_HERSHEY_SIMPLEX, - 0.6, [255, 255, 255], 1) + color = [random.randint(0, 255) for _ in range(3)] + if is_label: + label = str(id) + ":"+ names[cat] + (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1) + cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), color, -1) + cv2.putText(img, label, (x1, y1 - 5),cv2.FONT_HERSHEY_SIMPLEX, + 0.6, [255, 255, 255], 1) + cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) # cv2.circle(img, data, 6, color,-1) #centroid of box txt_str = "" if save_with_object_id: @@ -234,7 +236,7 @@ def detect(save_img=False): bbox_xyxy = tracked_dets[:,:4] identities = tracked_dets[:, 8] categories = tracked_dets[:, 4] - draw_boxes(im0, bbox_xyxy, identities, categories, names, save_with_object_id, txt_path) + draw_boxes(im0, bbox_xyxy, identities, categories, names, save_with_object_id, txt_path, is_label=True) else: #SORT should be updated even with no detections tracked_dets = sort_tracker.update() #........................................................