start_no_thread.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712
  1. import datetime
  2. import time
  3. import threading
  4. import queue
  5. import itertools
  6. from collections import defaultdict
  7. from statistics import mean
  8. import cv2
  9. import imutils
  10. import numpy as np
  11. import subprocess as sp
  12. import multiprocessing as mp
  13. import SharedArray as sa
  14. from scipy.spatial import distance as dist
  15. import tflite_runtime.interpreter as tflite
  16. from tflite_runtime.interpreter import load_delegate
  17. def load_labels(path, encoding='utf-8'):
  18. """Loads labels from file (with or without index numbers).
  19. Args:
  20. path: path to label file.
  21. encoding: label file encoding.
  22. Returns:
  23. Dictionary mapping indices to labels.
  24. """
  25. with open(path, 'r', encoding=encoding) as f:
  26. lines = f.readlines()
  27. if not lines:
  28. return {}
  29. if lines[0].split(' ', maxsplit=1)[0].isdigit():
  30. pairs = [line.split(' ', maxsplit=1) for line in lines]
  31. return {int(index): label.strip() for index, label in pairs}
  32. else:
  33. return {index: line.strip() for index, line in enumerate(lines)}
  34. def draw_box_with_label(frame, x_min, y_min, x_max, y_max, label, info, thickness=2, color=None, position='ul'):
  35. if color is None:
  36. color = (0,0,255)
  37. display_text = "{}: {}".format(label, info)
  38. cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, thickness)
  39. font_scale = 0.5
  40. font = cv2.FONT_HERSHEY_SIMPLEX
  41. # get the width and height of the text box
  42. size = cv2.getTextSize(display_text, font, fontScale=font_scale, thickness=2)
  43. text_width = size[0][0]
  44. text_height = size[0][1]
  45. line_height = text_height + size[1]
  46. # set the text start position
  47. if position == 'ul':
  48. text_offset_x = x_min
  49. text_offset_y = 0 if y_min < line_height else y_min - (line_height+8)
  50. elif position == 'ur':
  51. text_offset_x = x_max - (text_width+8)
  52. text_offset_y = 0 if y_min < line_height else y_min - (line_height+8)
  53. elif position == 'bl':
  54. text_offset_x = x_min
  55. text_offset_y = y_max
  56. elif position == 'br':
  57. text_offset_x = x_max - (text_width+8)
  58. text_offset_y = y_max
  59. # make the coords of the box with a small padding of two pixels
  60. textbox_coords = ((text_offset_x, text_offset_y), (text_offset_x + text_width + 2, text_offset_y + line_height))
  61. cv2.rectangle(frame, textbox_coords[0], textbox_coords[1], color, cv2.FILLED)
  62. cv2.putText(frame, display_text, (text_offset_x, text_offset_y + line_height - 3), font, fontScale=font_scale, color=(0, 0, 0), thickness=2)
  63. def calculate_region(frame_shape, xmin, ymin, xmax, ymax, multiplier=2):
  64. # size is larger than longest edge
  65. size = int(max(xmax-xmin, ymax-ymin)*multiplier)
  66. # if the size is too big to fit in the frame
  67. if size > min(frame_shape[0], frame_shape[1]):
  68. size = min(frame_shape[0], frame_shape[1])
  69. # x_offset is midpoint of bounding box minus half the size
  70. x_offset = int((xmax-xmin)/2.0+xmin-size/2.0)
  71. # if outside the image
  72. if x_offset < 0:
  73. x_offset = 0
  74. elif x_offset > (frame_shape[1]-size):
  75. x_offset = (frame_shape[1]-size)
  76. # y_offset is midpoint of bounding box minus half the size
  77. y_offset = int((ymax-ymin)/2.0+ymin-size/2.0)
  78. # if outside the image
  79. if y_offset < 0:
  80. y_offset = 0
  81. elif y_offset > (frame_shape[0]-size):
  82. y_offset = (frame_shape[0]-size)
  83. return (x_offset, y_offset, x_offset+size, y_offset+size)
  84. def intersection(box_a, box_b):
  85. return (
  86. max(box_a[0], box_b[0]),
  87. max(box_a[1], box_b[1]),
  88. min(box_a[2], box_b[2]),
  89. min(box_a[3], box_b[3])
  90. )
  91. def area(box):
  92. return (box[2]-box[0] + 1)*(box[3]-box[1] + 1)
  93. def intersection_over_union(box_a, box_b):
  94. # determine the (x, y)-coordinates of the intersection rectangle
  95. intersect = intersection(box_a, box_b)
  96. # compute the area of intersection rectangle
  97. inter_area = max(0, intersect[2] - intersect[0] + 1) * max(0, intersect[3] - intersect[1] + 1)
  98. if inter_area == 0:
  99. return 0.0
  100. # compute the area of both the prediction and ground-truth
  101. # rectangles
  102. box_a_area = (box_a[2] - box_a[0] + 1) * (box_a[3] - box_a[1] + 1)
  103. box_b_area = (box_b[2] - box_b[0] + 1) * (box_b[3] - box_b[1] + 1)
  104. # compute the intersection over union by taking the intersection
  105. # area and dividing it by the sum of prediction + ground-truth
  106. # areas - the interesection area
  107. iou = inter_area / float(box_a_area + box_b_area - inter_area)
  108. # return the intersection over union value
  109. return iou
  110. def clipped(obj, frame_shape):
  111. # if the object is within 5 pixels of the region border, and the region is not on the edge
  112. # consider the object to be clipped
  113. box = obj[2]
  114. region = obj[3]
  115. if ((region[0] > 5 and box[0]-region[0] <= 5) or
  116. (region[1] > 5 and box[1]-region[1] <= 5) or
  117. (frame_shape[1]-region[2] > 5 and region[2]-box[2] <= 5) or
  118. (frame_shape[0]-region[3] > 5 and region[3]-box[3] <= 5)):
  119. return True
  120. else:
  121. return False
  122. def filtered(obj):
  123. if obj[0] != 'person':
  124. return True
  125. return False
  126. def create_tensor_input(frame, region):
  127. cropped_frame = frame[region[1]:region[3], region[0]:region[2]]
  128. # Resize to 300x300 if needed
  129. if cropped_frame.shape != (300, 300, 3):
  130. # TODO: use Pillow-SIMD?
  131. cropped_frame = cv2.resize(cropped_frame, dsize=(300, 300), interpolation=cv2.INTER_LINEAR)
  132. # Expand dimensions since the model expects images to have shape: [1, 300, 300, 3]
  133. return np.expand_dims(cropped_frame, axis=0)
  134. class MotionDetector():
  135. # TODO: add motion masking
  136. def __init__(self, frame_shape, resize_factor=4):
  137. self.resize_factor = resize_factor
  138. self.motion_frame_size = (int(frame_shape[0]/resize_factor), int(frame_shape[1]/resize_factor))
  139. self.avg_frame = np.zeros(self.motion_frame_size, np.float)
  140. self.avg_delta = np.zeros(self.motion_frame_size, np.float)
  141. self.motion_frame_count = 0
  142. self.frame_counter = 0
  143. def detect(self, frame):
  144. motion_boxes = []
  145. # resize frame
  146. resized_frame = cv2.resize(frame, dsize=(self.motion_frame_size[1], self.motion_frame_size[0]), interpolation=cv2.INTER_LINEAR)
  147. # convert to grayscale
  148. gray = cv2.cvtColor(resized_frame, cv2.COLOR_BGR2GRAY)
  149. # it takes ~30 frames to establish a baseline
  150. # dont bother looking for motion
  151. if self.frame_counter < 30:
  152. self.frame_counter += 1
  153. else:
  154. # compare to average
  155. frameDelta = cv2.absdiff(gray, cv2.convertScaleAbs(self.avg_frame))
  156. # compute the average delta over the past few frames
  157. # the alpha value can be modified to configure how sensitive the motion detection is.
  158. # higher values mean the current frame impacts the delta a lot, and a single raindrop may
  159. # register as motion, too low and a fast moving person wont be detected as motion
  160. # this also assumes that a person is in the same location across more than a single frame
  161. cv2.accumulateWeighted(frameDelta, self.avg_delta, 0.2)
  162. # compute the threshold image for the current frame
  163. current_thresh = cv2.threshold(frameDelta, 25, 255, cv2.THRESH_BINARY)[1]
  164. # black out everything in the avg_delta where there isnt motion in the current frame
  165. avg_delta_image = cv2.convertScaleAbs(self.avg_delta)
  166. avg_delta_image[np.where(current_thresh==[0])] = [0]
  167. # then look for deltas above the threshold, but only in areas where there is a delta
  168. # in the current frame. this prevents deltas from previous frames from being included
  169. thresh = cv2.threshold(avg_delta_image, 25, 255, cv2.THRESH_BINARY)[1]
  170. # dilate the thresholded image to fill in holes, then find contours
  171. # on thresholded image
  172. thresh = cv2.dilate(thresh, None, iterations=2)
  173. cnts = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  174. cnts = imutils.grab_contours(cnts)
  175. # loop over the contours
  176. for c in cnts:
  177. # if the contour is big enough, count it as motion
  178. contour_area = cv2.contourArea(c)
  179. if contour_area > 100:
  180. # cv2.drawContours(resized_frame, [c], -1, (255,255,255), 2)
  181. x, y, w, h = cv2.boundingRect(c)
  182. motion_boxes.append((x*self.resize_factor, y*self.resize_factor, (x+w)*self.resize_factor, (y+h)*self.resize_factor))
  183. if len(motion_boxes) > 0:
  184. self.motion_frame_count += 1
  185. # TODO: this really depends on FPS
  186. if self.motion_frame_count >= 10:
  187. # only average in the current frame if the difference persists for at least 3 frames
  188. cv2.accumulateWeighted(gray, self.avg_frame, 0.2)
  189. else:
  190. # when no motion, just keep averaging the frames together
  191. cv2.accumulateWeighted(gray, self.avg_frame, 0.2)
  192. self.motion_frame_count = 0
  193. return motion_boxes
  194. class ObjectDetector():
  195. def __init__(self, model_file, label_file):
  196. self.labels = load_labels(label_file)
  197. edge_tpu_delegate = None
  198. try:
  199. edge_tpu_delegate = load_delegate('libedgetpu.so.1.0')
  200. except ValueError:
  201. print("No EdgeTPU detected. Falling back to CPU.")
  202. if edge_tpu_delegate is None:
  203. self.interpreter = tflite.Interpreter(
  204. model_path=model_file)
  205. else:
  206. self.interpreter = tflite.Interpreter(
  207. model_path=model_file,
  208. experimental_delegates=[edge_tpu_delegate])
  209. self.interpreter.allocate_tensors()
  210. self.tensor_input_details = self.interpreter.get_input_details()
  211. self.tensor_output_details = self.interpreter.get_output_details()
  212. def detect_raw(self, tensor_input):
  213. self.interpreter.set_tensor(self.tensor_input_details[0]['index'], tensor_input)
  214. self.interpreter.invoke()
  215. boxes = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[0]['index']))
  216. label_codes = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[1]['index']))
  217. scores = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[2]['index']))
  218. detections = np.zeros((20,6), np.float32)
  219. for i, score in enumerate(scores):
  220. detections[i] = [label_codes[i], score, boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]]
  221. return detections
  222. def detect(self, tensor_input, threshold=.4):
  223. self.interpreter.set_tensor(self.tensor_input_details[0]['index'], tensor_input)
  224. self.interpreter.invoke()
  225. boxes = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[0]['index']))
  226. label_codes = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[1]['index']))
  227. scores = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[2]['index']))
  228. detections = []
  229. for i, score in enumerate(scores):
  230. label = self.labels[int(label_codes[i])]
  231. if score < threshold:
  232. break
  233. detections.append((
  234. label,
  235. float(score),
  236. boxes[i]
  237. ))
  238. return detections
  239. class RemoteObjectDetector():
  240. def __init__(self, model, labels):
  241. self.labels = load_labels(labels)
  242. try:
  243. sa.delete("frame")
  244. except:
  245. pass
  246. try:
  247. sa.delete("detections")
  248. except:
  249. pass
  250. self.input_frame = sa.create("frame", shape=(1,300,300,3), dtype=np.uint8)
  251. self.detections = sa.create("detections", shape=(20,6), dtype=np.float32)
  252. self.detect_lock = mp.Lock()
  253. self.detect_ready = mp.Event()
  254. self.frame_ready = mp.Event()
  255. def run_detector(model, labels, detect_ready, frame_ready):
  256. object_detector = ObjectDetector(model, labels)
  257. input_frame = sa.attach("frame")
  258. detections = sa.attach("detections")
  259. while True:
  260. # signal that the process is ready to detect
  261. detect_ready.set()
  262. # wait until a frame is ready
  263. frame_ready.wait()
  264. # signal that the process is busy
  265. detect_ready.clear()
  266. frame_ready.clear()
  267. detections[:] = object_detector.detect_raw(input_frame)
  268. self.detect_process = mp.Process(target=run_detector, args=(model, labels, self.detect_ready, self.frame_ready))
  269. self.detect_process.daemon = True
  270. self.detect_process.start()
  271. def detect(self, tensor_input, threshold=.4):
  272. detections = []
  273. with self.detect_lock:
  274. self.input_frame[:] = tensor_input
  275. # signal that a frame is ready
  276. self.frame_ready.set()
  277. # wait until the detection process is finished,
  278. self.detect_ready.wait()
  279. for d in self.detections:
  280. if d[1] < threshold:
  281. break
  282. detections.append((
  283. self.labels[int(d[0])],
  284. float(d[1]),
  285. (d[2], d[3], d[4], d[5])
  286. ))
  287. return detections
  288. class ObjectTracker():
  289. def __init__(self, max_disappeared):
  290. self.tracked_objects = {}
  291. self.disappeared = {}
  292. self.max_disappeared = max_disappeared
  293. def register(self, index, frame_time, obj):
  294. id = f"{frame_time}-{index}"
  295. obj['id'] = id
  296. obj['frame_time'] = frame_time
  297. obj['top_score'] = obj['score']
  298. self.add_history(obj)
  299. self.tracked_objects[id] = obj
  300. self.disappeared[id] = 0
  301. def deregister(self, id):
  302. del self.tracked_objects[id]
  303. del self.disappeared[id]
  304. def update(self, id, new_obj):
  305. self.disappeared[id] = 0
  306. self.tracked_objects[id].update(new_obj)
  307. self.add_history(self.tracked_objects[id])
  308. if self.tracked_objects[id]['score'] > self.tracked_objects[id]['top_score']:
  309. self.tracked_objects[id]['top_score'] = self.tracked_objects[id]['score']
  310. def add_history(self, obj):
  311. entry = {
  312. 'score': obj['score'],
  313. 'box': obj['box'],
  314. 'region': obj['region'],
  315. 'centroid': obj['centroid'],
  316. 'frame_time': obj['frame_time']
  317. }
  318. if 'history' in obj:
  319. obj['history'].append(entry)
  320. else:
  321. obj['history'] = [entry]
  322. def match_and_update(self, frame_time, new_objects):
  323. if len(new_objects) == 0:
  324. for id in list(self.tracked_objects.keys()):
  325. if self.disappeared[id] >= self.max_disappeared:
  326. self.deregister(id)
  327. else:
  328. self.disappeared[id] += 1
  329. return
  330. # group by name
  331. new_object_groups = defaultdict(lambda: [])
  332. for obj in new_objects:
  333. new_object_groups[obj[0]].append({
  334. 'label': obj[0],
  335. 'score': obj[1],
  336. 'box': obj[2],
  337. 'region': obj[3]
  338. })
  339. # track objects for each label type
  340. for label, group in new_object_groups.items():
  341. current_objects = [o for o in self.tracked_objects.values() if o['label'] == label]
  342. current_ids = [o['id'] for o in current_objects]
  343. current_centroids = np.array([o['centroid'] for o in current_objects])
  344. # compute centroids of new objects
  345. for obj in group:
  346. centroid_x = int((obj['box'][0]+obj['box'][2]) / 2.0)
  347. centroid_y = int((obj['box'][1]+obj['box'][3]) / 2.0)
  348. obj['centroid'] = (centroid_x, centroid_y)
  349. if len(current_objects) == 0:
  350. for index, obj in enumerate(group):
  351. self.register(index, frame_time, obj)
  352. return
  353. new_centroids = np.array([o['centroid'] for o in group])
  354. # compute the distance between each pair of tracked
  355. # centroids and new centroids, respectively -- our
  356. # goal will be to match each new centroid to an existing
  357. # object centroid
  358. D = dist.cdist(current_centroids, new_centroids)
  359. # in order to perform this matching we must (1) find the
  360. # smallest value in each row and then (2) sort the row
  361. # indexes based on their minimum values so that the row
  362. # with the smallest value is at the *front* of the index
  363. # list
  364. rows = D.min(axis=1).argsort()
  365. # next, we perform a similar process on the columns by
  366. # finding the smallest value in each column and then
  367. # sorting using the previously computed row index list
  368. cols = D.argmin(axis=1)[rows]
  369. # in order to determine if we need to update, register,
  370. # or deregister an object we need to keep track of which
  371. # of the rows and column indexes we have already examined
  372. usedRows = set()
  373. usedCols = set()
  374. # loop over the combination of the (row, column) index
  375. # tuples
  376. for (row, col) in zip(rows, cols):
  377. # if we have already examined either the row or
  378. # column value before, ignore it
  379. if row in usedRows or col in usedCols:
  380. continue
  381. # otherwise, grab the object ID for the current row,
  382. # set its new centroid, and reset the disappeared
  383. # counter
  384. objectID = current_ids[row]
  385. self.update(objectID, group[col])
  386. # indicate that we have examined each of the row and
  387. # column indexes, respectively
  388. usedRows.add(row)
  389. usedCols.add(col)
  390. # compute the column index we have NOT yet examined
  391. unusedRows = set(range(0, D.shape[0])).difference(usedRows)
  392. unusedCols = set(range(0, D.shape[1])).difference(usedCols)
  393. # in the event that the number of object centroids is
  394. # equal or greater than the number of input centroids
  395. # we need to check and see if some of these objects have
  396. # potentially disappeared
  397. if D.shape[0] >= D.shape[1]:
  398. for row in unusedRows:
  399. id = current_ids[row]
  400. if self.disappeared[id] >= self.max_disappeared:
  401. self.deregister(id)
  402. else:
  403. self.disappeared[id] += 1
  404. # if the number of input centroids is greater
  405. # than the number of existing object centroids we need to
  406. # register each new input centroid as a trackable object
  407. else:
  408. for col in unusedCols:
  409. self.register(col, frame_time, group[col])
  410. def main():
  411. frames = 0
  412. # frame_queue = queue.Queue(maxsize=5)
  413. # frame_cache = {}
  414. frame_shape = (1080,1920,3)
  415. # frame_shape = (720,1280,3)
  416. frame_size = frame_shape[0]*frame_shape[1]*frame_shape[2]
  417. frame = np.zeros(frame_shape, np.uint8)
  418. motion_detector = MotionDetector(frame_shape, resize_factor=6)
  419. object_detector = ObjectDetector('/lab/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite', '/lab/labelmap.txt')
  420. # object_detector = RemoteObjectDetector('/lab/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite', '/lab/labelmap.txt')
  421. # object_detector = ObjectDetector('/lab/detect.tflite', '/lab/labelmap.txt')
  422. object_tracker = ObjectTracker(10)
  423. # f = open('/debug/input/back.rgb24', 'rb')
  424. # f = open('/debug/back.raw_video', 'rb')
  425. # f = open('/debug/ali-jake.raw_video', 'rb')
  426. # -hwaccel vaapi -hwaccel_device /dev/dri/renderD128 -hwaccel_output_format yuv420p -i output.mp4 -f rawvideo -pix_fmt rgb24 pipe:
  427. ffmpeg_cmd = (['ffmpeg'] +
  428. ['-hide_banner','-loglevel','panic'] +
  429. ['-hwaccel','vaapi','-hwaccel_device','/dev/dri/renderD129','-hwaccel_output_format','yuv420p'] +
  430. # ['-i', '/debug/input/output.mp4'] +
  431. ['-i', '/debug/back-ali-jake.mp4'] +
  432. ['-f','rawvideo','-pix_fmt','rgb24'] +
  433. ['pipe:'])
  434. print(" ".join(ffmpeg_cmd))
  435. ffmpeg_process = sp.Popen(ffmpeg_cmd, stdout = sp.PIPE, bufsize=frame_size)
  436. total_detections = 0
  437. start = datetime.datetime.now().timestamp()
  438. frame_times = []
  439. while True:
  440. start_frame = datetime.datetime.now().timestamp()
  441. frame_detections = 0
  442. frame_bytes = ffmpeg_process.stdout.read(frame_size)#f.read(frame_size)
  443. if not frame_bytes:
  444. break
  445. frame_time = datetime.datetime.now().timestamp()
  446. # Store frame in numpy array
  447. frame[:] = (np
  448. .frombuffer(frame_bytes, np.uint8)
  449. .reshape(frame_shape))
  450. frames += 1
  451. # look for motion
  452. motion_boxes = motion_detector.detect(frame)
  453. tracked_objects = object_tracker.tracked_objects.values()
  454. # merge areas of motion that intersect with a known tracked object into a single area to look at
  455. areas_of_interest = []
  456. used_motion_boxes = []
  457. for obj in tracked_objects:
  458. x_min, y_min, x_max, y_max = obj['box']
  459. for m_index, motion_box in enumerate(motion_boxes):
  460. if area(intersection(obj['box'], motion_box))/area(motion_box) > .5:
  461. used_motion_boxes.append(m_index)
  462. x_min = min(obj['box'][0], motion_box[0])
  463. y_min = min(obj['box'][1], motion_box[1])
  464. x_max = max(obj['box'][2], motion_box[2])
  465. y_max = max(obj['box'][3], motion_box[3])
  466. areas_of_interest.append((x_min, y_min, x_max, y_max))
  467. unused_motion_boxes = set(range(0, len(motion_boxes))).difference(used_motion_boxes)
  468. # compute motion regions
  469. motion_regions = [calculate_region(frame_shape, motion_boxes[i][0], motion_boxes[i][1], motion_boxes[i][2], motion_boxes[i][3], 1.2)
  470. for i in unused_motion_boxes]
  471. # compute tracked object regions
  472. object_regions = [calculate_region(frame_shape, a[0], a[1], a[2], a[3], 1.2)
  473. for a in areas_of_interest]
  474. # merge regions with high IOU
  475. merged_regions = motion_regions+object_regions
  476. while True:
  477. max_iou = 0.0
  478. max_indices = None
  479. region_indices = range(len(merged_regions))
  480. for a, b in itertools.combinations(region_indices, 2):
  481. iou = intersection_over_union(merged_regions[a], merged_regions[b])
  482. if iou > max_iou:
  483. max_iou = iou
  484. max_indices = (a, b)
  485. if max_iou > 0.1:
  486. a = merged_regions[max_indices[0]]
  487. b = merged_regions[max_indices[1]]
  488. merged_regions.append(calculate_region(frame_shape,
  489. min(a[0], b[0]),
  490. min(a[1], b[1]),
  491. max(a[2], b[2]),
  492. max(a[3], b[3]),
  493. 1
  494. ))
  495. del merged_regions[max(max_indices[0], max_indices[1])]
  496. del merged_regions[min(max_indices[0], max_indices[1])]
  497. else:
  498. break
  499. # resize regions and detect
  500. detections = []
  501. for region in merged_regions:
  502. tensor_input = create_tensor_input(frame, region)
  503. region_detections = object_detector.detect(tensor_input)
  504. frame_detections += 1
  505. for d in region_detections:
  506. if filtered(d):
  507. continue
  508. box = d[2]
  509. size = region[2]-region[0]
  510. x_min = int((box[1] * size) + region[0])
  511. y_min = int((box[0] * size) + region[1])
  512. x_max = int((box[3] * size) + region[0])
  513. y_max = int((box[2] * size) + region[1])
  514. detections.append((
  515. d[0],
  516. d[1],
  517. (x_min, y_min, x_max, y_max),
  518. region))
  519. #########
  520. # merge objects, check for clipped objects and look again up to N times
  521. #########
  522. refining = True
  523. refine_count = 0
  524. while refining and refine_count < 4:
  525. refining = False
  526. # group by name
  527. detected_object_groups = defaultdict(lambda: [])
  528. for detection in detections:
  529. detected_object_groups[detection[0]].append(detection)
  530. selected_objects = []
  531. for group in detected_object_groups.values():
  532. # apply non-maxima suppression to suppress weak, overlapping bounding boxes
  533. boxes = [(o[2][0], o[2][1], o[2][2]-o[2][0], o[2][3]-o[2][1])
  534. for o in group]
  535. confidences = [o[1] for o in group]
  536. idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
  537. for index in idxs:
  538. obj = group[index[0]]
  539. if clipped(obj, frame_shape): #obj['clipped']:
  540. box = obj[2]
  541. # calculate a new region that will hopefully get the entire object
  542. region = calculate_region(frame_shape,
  543. box[0], box[1],
  544. box[2], box[3])
  545. tensor_input = create_tensor_input(frame, region)
  546. # run detection on new region
  547. refined_detections = object_detector.detect(tensor_input)
  548. frame_detections += 1
  549. for d in refined_detections:
  550. if filtered(d):
  551. continue
  552. box = d[2]
  553. size = region[2]-region[0]
  554. x_min = int((box[1] * size) + region[0])
  555. y_min = int((box[0] * size) + region[1])
  556. x_max = int((box[3] * size) + region[0])
  557. y_max = int((box[2] * size) + region[1])
  558. selected_objects.append((
  559. d[0],
  560. d[1],
  561. (x_min, y_min, x_max, y_max),
  562. region))
  563. refining = True
  564. else:
  565. selected_objects.append(obj)
  566. # set the detections list to only include top, complete objects
  567. # and new detections
  568. detections = selected_objects
  569. if refining:
  570. refine_count += 1
  571. # now that we have refined our detections, we need to track objects
  572. object_tracker.match_and_update(frame_time, detections)
  573. total_detections += frame_detections
  574. frame_times.append(datetime.datetime.now().timestamp()-start_frame)
  575. # if (frames >= 700 and frames <= 1635) or (frames >= 2500):
  576. # if (frames >= 700 and frames <= 1000):
  577. if (frames >= 0):
  578. # row1 = cv2.hconcat([gray, cv2.convertScaleAbs(avg_frame)])
  579. # row2 = cv2.hconcat([frameDelta, thresh])
  580. # cv2.imwrite(f"/lab/debug/output/{frames}.jpg", cv2.vconcat([row1, row2]))
  581. # # cv2.imwrite(f"/lab/debug/output/resized-frame-{frames}.jpg", resized_frame)
  582. # for region in motion_regions:
  583. # cv2.rectangle(frame, (region[0], region[1]), (region[2], region[3]), (255,128,0), 2)
  584. # for region in object_regions:
  585. # cv2.rectangle(frame, (region[0], region[1]), (region[2], region[3]), (0,128,255), 2)
  586. for region in merged_regions:
  587. cv2.rectangle(frame, (region[0], region[1]), (region[2], region[3]), (0,255,0), 2)
  588. for box in motion_boxes:
  589. cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (255,0,0), 2)
  590. for detection in detections:
  591. box = detection[2]
  592. draw_box_with_label(frame, box[0], box[1], box[2], box[3], detection[0], f"{detection[1]*100}%")
  593. for obj in object_tracker.tracked_objects.values():
  594. box = obj['box']
  595. draw_box_with_label(frame, box[0], box[1], box[2], box[3], obj['label'], obj['id'], thickness=1, color=(0,0,255), position='bl')
  596. cv2.putText(frame, str(total_detections), (10, 10), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0, 0, 0), thickness=2)
  597. cv2.putText(frame, str(frame_detections), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0, 0, 0), thickness=2)
  598. cv2.imwrite(f"/lab/debug/output/frame-{frames}.jpg", frame)
  599. # break
  600. duration = datetime.datetime.now().timestamp()-start
  601. print(f"Processed {frames} frames for {duration:.2f} seconds and {(frames/duration):.2f} FPS.")
  602. print(f"Total detections: {total_detections}")
  603. print(f"Average frame processing time: {mean(frame_times)*1000:.2f}ms")
  604. if __name__ == '__main__':
  605. main()