Browse Source

Tracking seems to work OK initially

dev
sipp11 5 years ago
parent
commit
0b343a748a
  1. 57
      src/main.py
  2. 10
      src/utils.py
  3. 10
      src/yolo.py

57
src/main.py

@ -5,17 +5,18 @@ time python src/_detector.py --input ~/Desktop/5min.mp4 -l
""" """
# import the necessary packages # import the necessary packages
import numpy as np import time
import argparse import argparse
import pprint
import numpy as np
import imutils import imutils
import time
import cv2 import cv2
import pprint
from utils import ( from utils import (
check_if_inside_the_boxes, check_if_inside_the_boxes,
is_it_the_same_obj, is_it_the_same_obj,
box_distance, box_distance,
get_heading, get_heading,
get_avg_heading,
) )
pp = pprint.PrettyPrinter(indent=2) pp = pprint.PrettyPrinter(indent=2)
@ -25,6 +26,7 @@ OPENCV_OBJECT_TRACKERS = {"csrt": cv2.TrackerCSRT_create}
cv_trackers = cv2.MultiTracker_create() cv_trackers = cv2.MultiTracker_create()
trackers = [] trackers = []
finished = [] finished = []
W4A = {} # this stands for "wait for arrival [at ...]"
# construct the argument parse and parse the arguments # construct the argument parse and parse the arguments
ap = argparse.ArgumentParser() ap = argparse.ArgumentParser()
@ -72,7 +74,7 @@ ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()] ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
def detect_stuffs(net, frame): def detect_stuffs(frame):
# construct a blob from the input frame and then perform a forward # construct a blob from the input frame and then perform a forward
# pass of the YOLO object detector, giving us our bounding boxes # pass of the YOLO object detector, giving us our bounding boxes
# and associated probabilities # and associated probabilities
@ -228,11 +230,25 @@ while True:
if trackers[idx]["still"] > 30 or x < 5 or x > 1250: if trackers[idx]["still"] > 30 or x < 5 or x > 1250:
untracking.append(trackers[idx]) untracking.append(trackers[idx])
# DRAW on FRAME
tk = trackers[idx] tk = trackers[idx]
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2) # check if it's hit the first
# cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2) avg_heading = get_avg_heading(trackers[idx]["heading"])
color = [int(c) for c in COLORS[0]] print(f" ---> avg heading: ", avg_heading)
if (
tk["first"]["id"] == 4
and avg_heading == "N"
and (x > 140 or x + w > 175)
):
print("gone!")
_nid = f"id_5"
if _nid not in W4A:
W4A[_nid] = {"objects": []}
_expected_frame = 30 * 4 # at least 4 sec later
has_this = [_ for _ in W4A[_nid]["objects"] if _[0]["id"] == tk["id"]]
if not has_this:
W4A[_nid]["objects"].append((tk, _frame_count, _expected_frame))
untracking.append(trackers[idx])
print( print(
f"[{tk['id']}-{tk['type']}] (x,y)=({x},{y})" f"[{tk['id']}-{tk['type']}] (x,y)=({x},{y})"
f" | still #{tk['still']} | distance: " f" | still #{tk['still']} | distance: "
@ -240,6 +256,10 @@ while True:
) )
_htxt = ",".join(trackers[idx]["heading"]) _htxt = ",".join(trackers[idx]["heading"])
print(f" ------ heading: {_htxt}") print(f" ------ heading: {_htxt}")
# DRAW on FRAME
color = [int(c) for c in COLORS[0]]
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
cv2.putText( cv2.putText(
frame, frame,
f"{tk['id']} - {tk['type']}", f"{tk['id']} - {tk['type']}",
@ -252,7 +272,7 @@ while True:
# only detect once a sec # only detect once a sec
if _frame_count % 15 == 1: if _frame_count % 15 == 1:
idxs, boxes, confidences, classIDs, start, end = detect_stuffs(net, frame) idxs, boxes, confidences, classIDs, start, end = detect_stuffs(frame)
# loop over the indexes we are keeping # loop over the indexes we are keeping
for i in idxs.flatten(): for i in idxs.flatten():
# extract the bounding box coordinates # extract the bounding box coordinates
@ -277,6 +297,20 @@ while True:
break break
if not is_same: if not is_same:
gid = None
if found_at["id"] == 5:
print(f"FOUND AT 5 ON {_frame_count}")
_po = W4A["id_5"]["objects"]
_po = [_ for _ in _po if _frame_count > _[2]]
_po = sorted(_po, key=lambda kk: kk[2])
print(' ------ ', _po)
if _po:
gid = _po[0][0]["id"]
# remove this id out of next W4A
W4A["id_5"]["objects"] = [
_ for _ in W4A["id_5"]["objects"] if _[0]["id"] == gid
]
print(f" >> possibly this obj: ", _po)
# add tracker to this obj # add tracker to this obj
# create a new object tracker for the bounding box and add it # create a new object tracker for the bounding box and add it
# to our multi-object tracker # to our multi-object tracker
@ -286,7 +320,7 @@ while True:
bbox = (x, y, w, h) bbox = (x, y, w, h)
cv_trackers.add(_tracker, frame, bbox) cv_trackers.add(_tracker, frame, bbox)
t = { t = {
"id": tracker_counter, "id": tracker_counter if gid is None else gid,
"type": _class, "type": _class,
"curr_position": bbox, "curr_position": bbox,
"heading": [], "heading": [],
@ -295,7 +329,8 @@ while True:
"last_position": bbox, "last_position": bbox,
"still": 0, "still": 0,
} }
tracker_counter += 1 if gid is None:
tracker_counter += 1
trackers.append(t) trackers.append(t)
print(f"trackers ADD - now total #{len(trackers)}") print(f"trackers ADD - now total #{len(trackers)}")
pp.pprint(t) pp.pprint(t)

10
src/utils.py

@ -1,3 +1,4 @@
import collections
import math import math
# detecting area # detecting area
@ -64,6 +65,14 @@ AREAS = [
] ]
def get_avg_heading(headings):
latest = headings[:15]
chars = collections.Counter(''.join(latest)).most_common(10)
if chars:
return chars[0][0]
return None
def get_heading(x1, y1, x2, y2): def get_heading(x1, y1, x2, y2):
diff_x, diff_y = x2 - x1, y2 - y1 diff_x, diff_y = x2 - x1, y2 - y1
if diff_x > 0 and diff_y > 0: if diff_x > 0 and diff_y > 0:
@ -72,7 +81,6 @@ def get_heading(x1, y1, x2, y2):
return "N" return "N"
if diff_x > 0: if diff_x > 0:
return "NW" return "NW"
if diff_y > 0: if diff_y > 0:
return "SE" return "SE"
if diff_y == 0: if diff_y == 0:

10
src/yolo.py

@ -1,10 +0,0 @@
import cv2
import time
import numpy as np
def detect_stuffs(frame, net, ln, confidence, threshold, W, H):
return idxs, start, end
Loading…
Cancel
Save