speed-image-measurement-program

git clone git://git.codymlewis.com/speed-image-measurement-program.git
Log | Files | Refs | README

speedometer.py (3645B)


      1 #!/usr/bin/env python3
      2 # -*- coding: utf-8 -*-
      3 
      4 import cv2 as cv
      5 import numpy as np
      6 import dlib
      7 
      8 
      9 def norm(a, b, p=2):
     10     a = np.array(a)
     11     b = np.array(b)
     12     return np.power(np.sum(np.power(a - b, p)), 1/p)
     13 
     14 def kmeans(prev_frame, frame):
     15     centers = None
     16     flags = cv.KMEANS_RANDOM_CENTERS
     17     if prev_frame is not None:
     18         prev_frame = cv.cvtColor(prev_frame, cv.COLOR_BGR2GRAY)
     19         frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
     20         frame_diff = cv.subtract(frame, prev_frame)
     21         fx, fy = np.array(np.where(frame_diff > 6))
     22         frame_xy = np.stack([fy, fx], 1)
     23         if len(frame_xy) == 0:
     24             return None
     25         criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 1000, 2)
     26         _, _, centers = cv.kmeans(np.float32(frame_xy), 6, None, criteria, 5, flags)
     27         corners = set()
     28         for i, center in enumerate(centers):
     29             for j, other in enumerate(centers[i + 1:]):
     30                 dist = norm(center, other)
     31                 if dist < 120:
     32                     corners = corners.union({tuple(center), tuple(other)})
     33         if len(corners) == 4:
     34             min_x, min_y = None, None
     35             max_x, max_y = None, None
     36             for corner in corners:
     37                 if min_x is None or corner[0] < min_x:
     38                     min_x = int(corner[0])
     39                 if max_x is None or corner[0] > max_x:
     40                     max_x = int(corner[0])
     41                 if min_y is None or corner[1] < min_y:
     42                     min_y = int(corner[1])
     43                 if max_y is None or corner[1] > max_y:
     44                     max_y = int(corner[1])
     45             return (min_x, min_y, max_x - min_x, max_y - min_y)
     46     return None
     47 
     48 def calibrate(frame):
     49     r = dlib.get_frontal_face_detector()(frame)
     50     return (1.8 / r.pop().top()) if len(r) else None
     51 
     52 
     53 def track():
     54     cap = cv.VideoCapture(0)
     55     if cap.isOpened() == False:
     56         print("Error opening video stream or file")
     57     tracker = cv.TrackerTLD_create()
     58     init_bb = None
     59     prev_frame = None
     60     prev_box = None
     61     dist_unit = None
     62     while cap.isOpened():
     63         ret, frame = cap.read()
     64         if dist_unit is None:
     65             dist_unit = calibrate(frame)
     66         elif init_bb is not None:
     67             tick_count = cv.getTickCount()
     68             success, box = tracker.update(frame)
     69             fps = cv.getTickFrequency() / (cv.getTickCount() - tick_count)
     70             if success:
     71                 x, y, w, h = (int(i) for i in box)
     72                 frame = cv.rectangle(
     73                     frame,
     74                     (x, y),
     75                     (x + w, y + h),
     76                     (200, 70, 0)
     77                 )
     78                 if prev_box is not None:
     79                     dist = norm(box[0:1], prev_box[0:1])
     80                     speed = dist * dist_unit * fps
     81                     frame = cv.putText(
     82                         frame,
     83                         f"Speed: {speed:.3F}ms",
     84                         (50, 50),
     85                         cv.FONT_HERSHEY_SIMPLEX,
     86                         0.75,
     87                         (200, 70, 0),
     88                         2
     89                     )
     90             prev_box = box
     91         else:
     92             init_bb = kmeans(prev_frame, frame)
     93             if init_bb is not None:
     94                 tracker.init(frame, init_bb)
     95         cv.imshow('Frame', frame)
     96         prev_frame = frame
     97         if ret:
     98             key = cv.waitKey(1) & 0xFF
     99             if key == ord('q'):
    100                 break
    101         else:
    102             break
    103     cap.release()
    104     cv.destroyAllWindows()
    105 
    106 if __name__ == '__main__':
    107     print("Hello there, starting capture...")
    108     track()
    109     print("Bye.")