diff --git a/src/main.py b/src/main.py index 8a03ffc..1988692 100644 --- a/src/main.py +++ b/src/main.py @@ -10,8 +10,15 @@ import argparse import imutils import time import cv2 -from utils import check_if_inside_the_boxes, is_it_the_same_obj, box_distance +import pprint +from utils import ( + check_if_inside_the_boxes, + is_it_the_same_obj, + box_distance, + get_heading, +) +pp = pprint.PrettyPrinter(indent=2) # tracking OPENCV_OBJECT_TRACKERS = {"csrt": cv2.TrackerCSRT_create} # initialize OpenCV's special multi-object tracker @@ -194,7 +201,7 @@ while True: # loop over the bounding boxes and draw then on the frame if success: obj_cnt = len(boxes) - print(f'obj_cnt: ', obj_cnt, boxes) + print(f"obj_cnt: ", obj_cnt, boxes) for idx in range(obj_cnt): box = boxes[idx] (x, y, w, h) = [int(v) for v in box] @@ -205,12 +212,19 @@ while True: trackers[idx]["distance"] = curr_distance trackers[idx]["curr_position"] = [int(v) for v in box] + _heading = get_heading(_x, _y, x, y) + STILL_DIST_PX = 2 if last_distance < STILL_DIST_PX and curr_distance < STILL_DIST_PX: trackers[idx]["still"] += 1 else: trackers[idx]["still"] = 0 + if trackers[idx]["still"] == 0: + trackers[idx]["heading"] = [_heading] + trackers[idx]["heading"] + + trackers[idx]["heading"] = trackers[idx]["heading"][:30] + if trackers[idx]["still"] > 30 or x < 5 or x > 1250: untracking.append(trackers[idx]) @@ -220,8 +234,12 @@ while True: # cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2) color = [int(c) for c in COLORS[0]] print( - f"{tk['id']} - {tk['type']} - centroid: {x}, {y} - distance: [stl:{tk['still']}] {last_distance:.3f} -> {curr_distance:.3f}" + f"[{tk['id']}-{tk['type']}] (x,y)=({x},{y})" + f" | still #{tk['still']} | distance: " + f"{last_distance:.3f} -> {curr_distance:.3f}" ) + _htxt = ",".join(trackers[idx]["heading"]) + print(f" ------ heading: {_htxt}") cv2.putText( frame, f"{tk['id']} - {tk['type']}", @@ -242,8 +260,8 @@ while True: (w, h) = (boxes[i][2], boxes[i][3]) _class = LABELS[classIDs[i]] - _good = check_if_inside_the_boxes(x, y, w, h, _class) - if not _good: + found_at = check_if_inside_the_boxes(x, y, w, h, _class) + if not found_at: continue # (1) check whether it's the same object as one in trackers @@ -271,15 +289,17 @@ while True: "id": tracker_counter, "type": _class, "curr_position": bbox, - "direction": "", + "heading": [], + "first": found_at, "distance": -1, "last_position": bbox, "still": 0, } tracker_counter += 1 trackers.append(t) - print(f'trackers ADD - now total #{len(trackers)}') - print(f" i -> {i} ({x},{y}), {w},{h} ({x + w},{y + h})") + print(f"trackers ADD - now total #{len(trackers)}") + pp.pprint(t) + # print(f" i -> {i} ({x},{y}), {w},{h} ({x + w},{y + h})") _what = ",".join([LABELS[c] for c in classIDs]) print(f"[{_frame_count:08d}] :: {_what}") diff --git a/src/utils.py b/src/utils.py index 494ce90..1460201 100755 --- a/src/utils.py +++ b/src/utils.py @@ -6,9 +6,20 @@ AREAS = [ ("id", 1), ("area", ((0, 40), (12, 129))), ("target", ["car", "bus", "motorbike"]), + ("next", [6]), + ], + [ + ("id", 2), + ("area", ((85, 0), (222, 74))), + ("target", ["person", "bicycle"]), + ("next", [7]), + ], + [ + ("id", 3), + ("area", ((38, 340), (99, 482))), + ("target", ["person", "wheelchair"]), + ("next", [5]), ], - [("id", 2), ("area", ((85, 0), (222, 74))), ("target", ["person", "bicycle"])], - [("id", 3), ("area", ((38, 340), (99, 482))), ("target", ["person", "wheelchair"])], [ ("id", 4), ("area", ((106, 310), (164, 461))), @@ -18,35 +29,57 @@ AREAS = [ ("id", 5), ("area", ((286, 230), (441, 346))), ("target", ["person", "wheelchair"]), + ("next", [8]), ], [ ("id", 6), ("area", ((421, 190), (555, 304))), ("target", ["car", "bus", "motorbike"]), + ("next", []), ], [ ("id", 7), ("area", ((555, 170), (720, 295))), ("target", ["person", "wheelchair", "bicycle"]), + ("next", [4]), ], [ ("id", 8), ("area", ((877, 224), (947, 334))), ("target", ["person", "wheelchair"]), + ("next", [7, 9, 5]), ], [ ("id", 9), ("area", ((1047, 229), (112, 338))), ("target", ["person", "wheelchair"]), + ("next", [8]), ], [ ("id", 10), ("area", ((1158, 200), (1230, 307))), ("target", ["person", "wheelchair"]), + ("next", [9]), ], ] +def get_heading(x1, y1, x2, y2): + diff_x, diff_y = x2 - x1, y2 - y1 + if diff_x > 0 and diff_y > 0: + return "NE" + if diff_x > 0 and diff_y == 0: + return "N" + if diff_x > 0: + return "NW" + + if diff_y > 0: + return "SE" + if diff_y == 0: + return "S" + return "SW" + + def box_distance(pos1, pos2): x1, y1, w1, h1 = pos1 x2, y2, w2, h2 = pos2 @@ -75,7 +108,8 @@ def check_if_inside_the_boxes(x, y, w, h, _type): # if diff_x < box_w if is_inside: print("INSIDE!! this -> ", box) - return is_inside + return box + return False def is_it_the_same_obj(x1, y1, w1, h1, i1, j1, w2, h2, **kwargs):