Part 1. Object Detector

We need an object detector to detect the object so we can use the position in object tracking.

We use a pre-trained Yolov3 on MS coco to save time. You can select your own detector, just get the position.

Preparation: install required modules

Be sure to restart the runtime if instructed

In [1]:
!pip3 install imageai
!pip3 install keras
!pip3 install tensorflow
Collecting imageai
  Downloading imageai-2.1.6-py3-none-any.whl (160 kB)
     |████████████████████████████████| 160 kB 5.1 MB/s 
Collecting matplotlib==3.3.2
  Downloading matplotlib-3.3.2-cp37-cp37m-manylinux1_x86_64.whl (11.6 MB)
     |████████████████████████████████| 11.6 MB 49.5 MB/s 
Collecting pillow==7.0.0
  Downloading Pillow-7.0.0-cp37-cp37m-manylinux1_x86_64.whl (2.1 MB)
     |████████████████████████████████| 2.1 MB 54.4 MB/s 
Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (from imageai) (4.1.2.30)
Collecting keras-resnet==0.2.0
  Downloading keras-resnet-0.2.0.tar.gz (9.3 kB)
Collecting h5py==2.10.0
  Downloading h5py-2.10.0-cp37-cp37m-manylinux1_x86_64.whl (2.9 MB)
     |████████████████████████████████| 2.9 MB 52.9 MB/s 
Requirement already satisfied: scipy==1.4.1 in /usr/local/lib/python3.7/dist-packages (from imageai) (1.4.1)
Collecting numpy==1.19.3
  Downloading numpy-1.19.3-cp37-cp37m-manylinux2010_x86_64.whl (14.9 MB)
     |████████████████████████████████| 14.9 MB 57.5 MB/s 
Collecting keras==2.4.3
  Downloading Keras-2.4.3-py2.py3-none-any.whl (36 kB)
Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from h5py==2.10.0->imageai) (1.15.0)
Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from keras==2.4.3->imageai) (3.13)
Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.2->imageai) (2.8.2)
Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.2->imageai) (1.4.2)
Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.3 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.2->imageai) (3.0.8)
Requirement already satisfied: certifi>=2020.06.20 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.2->imageai) (2021.10.8)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.2->imageai) (0.11.0)
Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from kiwisolver>=1.0.1->matplotlib==3.3.2->imageai) (4.1.1)
Building wheels for collected packages: keras-resnet
  Building wheel for keras-resnet (setup.py) ... done
  Created wheel for keras-resnet: filename=keras_resnet-0.2.0-py2.py3-none-any.whl size=20486 sha256=a6cd28f8924a9efbf337b823a0142275fb3623a63f091f9422f4594f5cd88a35
  Stored in directory: /root/.cache/pip/wheels/bd/ef/06/5d65f696360436c3a423020c4b7fd8c558c09ef264a0e6c575
Successfully built keras-resnet
Installing collected packages: numpy, h5py, pillow, keras, matplotlib, keras-resnet, imageai
  Attempting uninstall: numpy
    Found existing installation: numpy 1.21.6
    Uninstalling numpy-1.21.6:
      Successfully uninstalled numpy-1.21.6
  Attempting uninstall: h5py
    Found existing installation: h5py 3.1.0
    Uninstalling h5py-3.1.0:
      Successfully uninstalled h5py-3.1.0
  Attempting uninstall: pillow
    Found existing installation: Pillow 7.1.2
    Uninstalling Pillow-7.1.2:
      Successfully uninstalled Pillow-7.1.2
  Attempting uninstall: keras
    Found existing installation: keras 2.8.0
    Uninstalling keras-2.8.0:
      Successfully uninstalled keras-2.8.0
  Attempting uninstall: matplotlib
    Found existing installation: matplotlib 3.2.2
    Uninstalling matplotlib-3.2.2:
      Successfully uninstalled matplotlib-3.2.2
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
tensorflow 2.8.0 requires tf-estimator-nightly==2.8.0.dev2021122109, which is not installed.
tensorflow 2.8.0 requires keras<2.9,>=2.8.0rc0, but you have keras 2.4.3 which is incompatible.
tensorflow 2.8.0 requires numpy>=1.20, but you have numpy 1.19.3 which is incompatible.
datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.
bokeh 2.3.3 requires pillow>=7.1.0, but you have pillow 7.0.0 which is incompatible.
albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.
Successfully installed h5py-2.10.0 imageai-2.1.6 keras-2.4.3 keras-resnet-0.2.0 matplotlib-3.3.2 numpy-1.19.3 pillow-7.0.0
Requirement already satisfied: keras in /usr/local/lib/python3.7/dist-packages (2.4.3)
Requirement already satisfied: h5py in /usr/local/lib/python3.7/dist-packages (from keras) (2.10.0)
Requirement already satisfied: numpy>=1.9.1 in /usr/local/lib/python3.7/dist-packages (from keras) (1.19.3)
Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from keras) (3.13)
Requirement already satisfied: scipy>=0.14 in /usr/local/lib/python3.7/dist-packages (from keras) (1.4.1)
Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from h5py->keras) (1.15.0)
Requirement already satisfied: tensorflow in /usr/local/lib/python3.7/dist-packages (2.8.0)
Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.24.0)
Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from tensorflow) (57.4.0)
Requirement already satisfied: absl-py>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.0.0)
Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.15.0)
Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.2.0)
Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.6.3)
Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.1.0)
Requirement already satisfied: gast>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (0.5.3)
Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.14.0)
Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.10.0)
Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.1.2)
Requirement already satisfied: libclang>=9.0.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (13.0.0)
Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.3.0)
Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (3.17.3)
Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (4.1.1)
Collecting tf-estimator-nightly==2.8.0.dev2021122109
  Downloading tf_estimator_nightly-2.8.0.dev2021122109-py2.py3-none-any.whl (462 kB)
     |████████████████████████████████| 462 kB 5.0 MB/s 
Collecting numpy>=1.20
  Downloading numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (15.7 MB)
     |████████████████████████████████| 15.7 MB 65.3 MB/s 
Requirement already satisfied: tensorboard<2.9,>=2.8 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.8.0)
Collecting keras<2.9,>=2.8.0rc0
  Downloading keras-2.8.0-py2.py3-none-any.whl (1.4 MB)
     |████████████████████████████████| 1.4 MB 56.8 MB/s 
Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (1.44.0)
Requirement already satisfied: flatbuffers>=1.12 in /usr/local/lib/python3.7/dist-packages (from tensorflow) (2.0)
Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.7/dist-packages (from astunparse>=1.6.0->tensorflow) (0.37.1)
Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (3.3.6)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (0.4.6)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (1.8.1)
Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (2.23.0)
Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (1.0.1)
Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (1.35.0)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow) (0.6.1)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow) (0.2.8)
Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow) (4.8)
Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow) (4.2.4)
Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow) (1.3.1)
Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow) (4.11.3)
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow) (3.8.0)
Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow) (0.4.8)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow) (2021.10.8)
Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow) (3.0.4)
Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow) (1.24.3)
Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow) (2.10)
Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow) (3.2.0)
Installing collected packages: numpy, tf-estimator-nightly, keras
  Attempting uninstall: numpy
    Found existing installation: numpy 1.19.3
    Uninstalling numpy-1.19.3:
      Successfully uninstalled numpy-1.19.3
  Attempting uninstall: keras
    Found existing installation: Keras 2.4.3
    Uninstalling Keras-2.4.3:
      Successfully uninstalled Keras-2.4.3
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
imageai 2.1.6 requires keras==2.4.3, but you have keras 2.8.0 which is incompatible.
imageai 2.1.6 requires numpy==1.19.3, but you have numpy 1.21.6 which is incompatible.
datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.
bokeh 2.3.3 requires pillow>=7.1.0, but you have pillow 7.0.0 which is incompatible.
albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.
Successfully installed keras-2.8.0 numpy-1.21.6 tf-estimator-nightly-2.8.0.dev2021122109

Download pretrained model

If all links below expired, please download the pre-trained model from imageai's official website: https://imageai.readthedocs.io/en/latest/detection/index.html

In [1]:
!wget https://dl.acytoo.com/yolo.h5 # if this site is down, use https://play.acytoo.com/yolo.h5 instead
--2022-04-23 06:15:26--  https://dl.acytoo.com/yolo.h5
Resolving dl.acytoo.com (dl.acytoo.com)... 129.158.55.57
Connecting to dl.acytoo.com (dl.acytoo.com)|129.158.55.57|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 248686624 (237M) [application/x-hdf]
Saving to: ‘yolo.h5’

yolo.h5             100%[===================>] 237.17M   123MB/s    in 1.9s    

2022-04-23 06:15:28 (123 MB/s) - ‘yolo.h5’ saved [248686624/248686624]

Define a general detector model: yolov3, trained on ms coco dataset.

Yolo(You Only Look Once) is one stage concolutional neural network for peforming object detection in real time. Compared to two stage CNNs, Yolo is faster and consumes less computing resources, but it's precision is lower than, it's a Speed-Accuracy Trade-Off.

Imageai is a python module that has many pre-trained object detection model.

For detailed information about imageai, please refer to https://imageai.readthedocs.io/en/latest/detection/index.html

In [2]:
from imageai.Detection import ObjectDetection

detector = ObjectDetection()
detector.setModelTypeAsYOLOv3()

Load pretrained model

In [3]:
detector.setModelPath("yolo.h5")
detector.loadModel()

Define ball detector from general detector

Yolo trained on ms coco can detect more than 170 kinds of object, we only need ball here, so after perform yolo on the input image, we select the objects that detedted as ball.

Ball detector, input is image, return probility and location

In [4]:
def ball_detector(img, threshold=30):
  """ 
  :param img: image matrix
  :param threshold: min probability that one object is detected as ball
  :return: [probability], [location]
  """
  img_ori_dec, detections = detector.detectObjectsFromImage(input_image=img, input_type = "array",
                                                            output_type="array", minimum_percentage_probability=threshold)
  prob = []
  loc = []
  for item in detections:
    if item['name'] == 'sports ball':
      prob.append(item['percentage_probability'])
      loc.append(item['box_points'])
  return prob, loc
In [5]:
%matplotlib inline
import matplotlib.pyplot as plt
import cv2

Download two videos

In [6]:
!wget https://dl.acytoo.com/ball.mp4
!wget https://dl.acytoo.com/multiObject.avi
--2022-04-23 06:16:01--  https://dl.acytoo.com/ball.mp4
Resolving dl.acytoo.com (dl.acytoo.com)... 129.158.55.57
Connecting to dl.acytoo.com (dl.acytoo.com)|129.158.55.57|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 1368142 (1.3M) [video/mp4]
Saving to: ‘ball.mp4’

ball.mp4            100%[===================>]   1.30M  --.-KB/s    in 0.08s   

2022-04-23 06:16:01 (15.4 MB/s) - ‘ball.mp4’ saved [1368142/1368142]

--2022-04-23 06:16:01--  https://dl.acytoo.com/multiObject.avi
Resolving dl.acytoo.com (dl.acytoo.com)... 129.158.55.57
Connecting to dl.acytoo.com (dl.acytoo.com)|129.158.55.57|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 1769608 (1.7M) [video/x-msvideo]
Saving to: ‘multiObject.avi’

multiObject.avi     100%[===================>]   1.69M  --.-KB/s    in 0.09s   

2022-04-23 06:16:01 (19.7 MB/s) - ‘multiObject.avi’ saved [1769608/1769608]

Part 2: Tracking

Get centroid of object

In above detector, our detector result is denoted as left top(x1, y1) and right bottom(x2, y2) vertexes, in order to apply kalman filter, we need to find the centroid of object(the ball). The centroid of ball can be calculated as $$ x_c = x_1 + \frac{(x_2 - x_1)}{2}\\ y_c = y_2 + \frac{(y_1 - y_2)}{2}\\ $$

In [7]:
def get_centroid(loc):
  if loc is None:
    return None
  return int(loc[0] + (loc[2] - loc[0]) / 2), int(loc[3] + (loc[1] - loc[3]) / 2)

Kalman filter

In [8]:
import numpy as np
from scipy.optimize import linear_sum_assignment # to assign detectors to objects
In [9]:
class KalmanFilter():
  def __init__(self, dt=0.1, u_x=1, u_y=1,
               std_acc=1, x_std_meas=0.1, y_std_meas=0.1):
    """
    :param dt: sampling time (time for 1 cycle)
    :param u_x: acceleration in x-direction
    :param u_y: acceleration in y-direction
    :param std_acc: process noise magnitude
    :param x_std_meas: standard deviation of the measurement in x-direction
    :param y_std_meas: standard deviation of the measurement in y-direction
    """

    self.dt = dt                              # Define sampling time
    self.u = np.matrix([[u_x],[u_y]])         # Define the  control input variables
    self.x = np.matrix([[0], [0], [0], [0]])  # Initial State

    # Define the State Transition Matrix A
    self.A = np.matrix([[1, 0, self.dt, 0],
                        [0, 1, 0, self.dt],
                        [0, 0, 1, 0],
                        [0, 0, 0, 1]])

    # Define the Control Input Matrix B
    self.B = np.matrix([[(self.dt**2)/2, 0],
                        [0,(self.dt**2)/2],
                        [self.dt,0],
                        [0,self.dt]])

    # Define Measurement Mapping Matrix
    self.H = np.matrix([[1, 0, 0, 0],
                        [0, 1, 0, 0]])

    #Initial Process Noise Covariance
    self.Q = np.matrix([[(self.dt**4)/4, 0, (self.dt**3)/2, 0],
                        [0, (self.dt**4)/4, 0, (self.dt**3)/2],
                        [(self.dt**3)/2, 0, self.dt**2, 0],
                        [0, (self.dt**3)/2, 0, self.dt**2]]) * std_acc**2

    #Initial Measurement Noise Covariance
    self.R = np.matrix([[x_std_meas**2,0],
                       [0, y_std_meas**2]])

    #Initial Covariance Matrix
    self.P = np.eye(self.A.shape[1])

  def predict(self):
    """
    :return: numpy matrix
    """
    # Update time state
    self.x = np.dot(self.A, self.x) + np.dot(self.B, self.u)
    # Calculate error covariance
    self.P = np.dot(np.dot(self.A, self.P), self.A.T) + self.Q
    return self.x[0:2]

  def update(self, z):
    """
    :param z: np.array
    :return: numpy matrix
    """
    # S = H*P*H'+R
    S = np.dot(self.H, np.dot(self.P, self.H.T)) + self.R
    # Calculate the Kalman Gain
    # K = P * H'* inv(H*P*H'+R)
    K = np.dot(np.dot(self.P, self.H.T), np.linalg.inv(S))
    self.x = np.round(self.x + np.dot(K, (z - np.dot(self.H, self.x))))
    I = np.eye(self.H.shape[1])
    # Update error covariance matrix
    self.P = (I - (K * self.H)) * self.P
    return self.x[0:2]

Green square is measurement, and red dot is prediction

In [10]:
def predict_ball(src, dest):
  """
  :param src: source video for detection and predicton
  :param dest: output video after processing

  detection: gree square box
  prediction: red dot
  """
  COLOR_RED = (255, 0, 0)
  COLOR_GREEN = (0, 255, 0)

  ############################ init kalman filter ##############################
  kalman = KalmanFilter()
  # set first point, the first location is VERY important
  kalman.update(np.array([[906], [311]], dtype=np.float32))


  vid = cv2.VideoCapture(src)
  if not vid.isOpened():
    print("File open error, please download video")

  ############################ params for saving video #########################
  fps = int(vid.get(cv2.CAP_PROP_FPS))              # fps
  height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))  # height
  width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))    # width
  size = (width, height)
  fourcc = cv2.VideoWriter_fourcc(*'MJPG')
  out = cv2.VideoWriter(dest, fourcc, fps, size)

  current_measurement = None
  current_prediction = (905, 311)

  while vid.isOpened():                             # read all video frames
    ret, frame = vid.read()
    if not ret:
      break

    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # convert bgr to rgb
    _, loc = ball_detector(frame)

    lpx, lpy = int(current_prediction[0]), int(current_prediction[1])
    frame = cv2.circle(frame, (lpx, lpy), 0, COLOR_RED, 20) # plot prediction dot

    if len(loc) > 0: # detection -> draw bounding box, update kalman filter
      frame = cv2.rectangle(frame, (loc[0][0], loc[0][1]), (loc[0][2], loc[0][3]), COLOR_GREEN, 2)
      # update kalman
      current_measurement = np.array(get_centroid(loc[0]), dtype=np.float32).reshape(2, 1)
      kalman.update(current_measurement)

    (x, y) = kalman.predict()
    current_prediction = (int(x), int(y))

    plt.imshow(frame, interpolation='nearest')
    plt.show()
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    out.write(frame)

  vid.release() # release video object
  out.release()
In [11]:
sin = 'ball.mp4'
sin_out = 'sin_track.avi'
print("processing sin")
predict_ball(sin, sin_out)
processing sin