diff --git a/asone/asone.py b/asone/asone.py index 93ba915..6e93fe6 100644 --- a/asone/asone.py +++ b/asone/asone.py @@ -31,18 +31,19 @@ def get_tracker(self, tracker: int): use_cuda=self.use_cuda) return tracker - def track_video(self, video_path, output_dir='results', save_result=True, display=False): + def track_video(self, video_path, output_dir='results', save_result=True, display=False, draw_trails=False): output_filename = os.path.basename(video_path) for (bbox_details, frame_details) in self._start_tracking(video_path, output_filename, output_dir=output_dir, save_result=save_result, - display=display): + display=display, + draw_trails=draw_trails): # yeild bbox_details, frame_details to main script yield bbox_details, frame_details - def track_webcam(self, cam_id=0, output_dir='results', save_result=False, display=True): + def track_webcam(self, cam_id=0, output_dir='results', save_result=False, display=True, draw_trails=False): output_filename = 'results.mp4' @@ -51,11 +52,13 @@ def track_webcam(self, cam_id=0, output_dir='results', save_result=False, displa output_dir=output_dir, fps=29, save_result=save_result, - display=display): + display=display, + draw_trails=draw_trails): # yeild bbox_details, frame_details to main script yield bbox_details, frame_details - def _start_tracking(self, stream_path, filename, fps=None, output_dir='results', save_result=True, display=False): + def _start_tracking(self, stream_path, filename, fps=None, output_dir='results', + save_result=True, display=False, draw_trails=False): cap = cv2.VideoCapture(stream_path) width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) @@ -98,7 +101,7 @@ def _start_tracking(self, stream_path, filename, fps=None, output_dir='results' 'frame {}/{} ({:.2f} ms)'.format(frame_id, int(frame_count), elapsed_time * 1000), ) - im0 = utils.draw_boxes(im0, bboxes_xyxy, identities=ids) + im0 = utils.draw_boxes(im0, bboxes_xyxy, identities=ids, draw_trails=draw_trails) currTime = time.time() fps = 1 / (currTime - prevTime) diff --git a/asone/utils/draw.py b/asone/utils/draw.py index e6cc65c..71043e8 100644 --- a/asone/utils/draw.py +++ b/asone/utils/draw.py @@ -3,8 +3,10 @@ import numpy as np from asone.utils import compute_color_for_labels from asone.utils import get_names +from collections import deque names = get_names() +data_deque = {} def draw_ui_box(x, img, id, color=None, label=None, line_thickness=None): @@ -42,9 +44,7 @@ def draw_ui_box(x, img, id, color=None, label=None, line_thickness=None): def draw_border(img, pt1, pt2, color, thickness, r, d): x1, y1 = pt1 x2, y2 = pt2 - # Top left - cv2.line(img, (x1 + r, y1), (x1 + r + d, y1), color, thickness) - cv2.line(img, (x1, y1 + r), (x1, y1 + r + d), color, thickness) + # Top leftfrom collections import deque (x1, y1 + r + d), color, thickness) cv2.ellipse(img, (x1 + r, y1 + r), (r, r), 180, 0, 90, color, thickness) # Top right @@ -71,10 +71,16 @@ def draw_border(img, pt1, pt2, color, thickness, r, d): return img -def draw_boxes(img, bbox_xyxy, class_ids=None, identities=None, offset=(0, 0)): +def draw_boxes(img, bbox_xyxy, class_ids=None, identities=None, draw_trails=False, offset=(0, 0)): # cv2.line(img, line2[0], line2[1], (0,200,0), 3) height, width, _ = img.shape + # remove tracked point from buffer if object is lost + if draw_trails: + for key in list(data_deque): + if key not in identities: + data_deque.pop(key) + for i, box in enumerate(bbox_xyxy): x1, y1, x2, y2 = [int(i) for i in box] x1 += offset[0] @@ -85,18 +91,46 @@ def draw_boxes(img, bbox_xyxy, class_ids=None, identities=None, offset=(0, 0)): # box_area = (x2-x1) * (y2-y1) box_height = (y2-y1) - # code to find center of bottom edge - center = (int((x2+x1) / 2), int((y2+y2)/2)) - # get ID of object id = int(identities[i]) if identities is not None else None color = compute_color_for_labels(0) + label = None if class_ids is not None: color = compute_color_for_labels(int(class_ids[i])) obj_name = names[int(class_ids[i])] label = '%s' % (obj_name) draw_ui_box(box, img, id=id, label=label, color=color, line_thickness=2) + + # Draw trails + + # code to find center of bottom edge + center = (int((x2+x1) / 2), int((y2+y2)/2)) + + if draw_trails: + # create new buffer for new object + if id not in data_deque: + data_deque[id] = deque(maxlen= 64) + + data_deque[id].appendleft(center) + drawtrails(data_deque, id, color, img) + + + return img + +def drawtrails(data_deque, id, color, img): + # draw trail + for i in range(1, len(data_deque[id])): + # check if on buffer value is none + if data_deque[id][i - 1] is None or data_deque[id][i] is None: + continue + + # generate dynamic thickness of trails + thickness = int(np.sqrt(64 / float(i + i)) * 1.5) + + # draw trails + cv2.line(img, data_deque[id][i - 1], data_deque[id][i], color, thickness) + diff --git a/main.py b/main.py index 9867199..03c2044 100644 --- a/main.py +++ b/main.py @@ -9,12 +9,11 @@ def main(args): detector=asone.YOLOX_DARKNET_PYTORCH, use_cuda=args.use_cuda ) - for bbox_details, frame_details in dt_obj.track_video( - args.video_path, - output_dir=args.output_dir, - save_result=args.save_result, - display=args.display - ): + for bbox_details, frame_details in dt_obj.track_video(args.video_path, + output_dir=args.output_dir, + save_result=args.save_result, + display=args.display, + draw_trails=args.draw_trails): # bbox_details = bbox_xyxy, ids, scores, class_ids # frame_details = frame, frame_no, fps @@ -31,6 +30,7 @@ def main(args): parser.add_argument('--no_display', default=True, action='store_false', dest='display', help='if provided the results will not be displayed on screen') parser.add_argument('--output_dir', default='data/results', help='Path to output directory') + parser.add_argument('--draw_trails', default=False, help='if provided object motion trails will be drawn.') args = parser.parse_args() diff --git a/setup.py b/setup.py index 6adea87..46ac9a9 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ URL = 'https://github.com/axcelerateai/asone' DOWNLOAD_URL = URL -VERSION = '0.1.2.dev8' +VERSION = '0.1.2.dev10' with open('README.md') as f: long_description = f.read()