Welcome to the object detection inference walkthrough! This notebook will walk you step by step through the process of using a pre-trained model to detect objects in an image. Make sure to follow the installation instructions before you start.
from __future__ import division
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
import math
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
# Import everything needed to edit/save/watch video clips
# from moviepy.editor import VideoFileClip
# from IPython.display import HTML
# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")
from object_detection.utils import ops as utils_ops
import cv2
Here are the imports from the object detection module.
from utils import label_map_util
from utils import visualization_utils as vis_util
Any model exported using the export_inference_graph.py
tool can be loaded here simply by changing PATH_TO_CKPT
to point to a new .pb file.
By default we use an "SSD with Mobilenet" model here. See the detection model zoo for a list of other models that can be run out-of-the-box with varying speeds and accuracies.
# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = '/home/priya/Documents/AI_Apps/soccer_project/exported_graphs_inception' + '/frozen_inference_graph.pb'
# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = '/home/priya/Documents/AI_Apps/soccer_project/data/' + 'soccer_label_map.pbtxt'
NUM_CLASSES = 2
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
Label maps map indices to category names, so that when our convolution network predicts 5
, we know that this corresponds to airplane
. Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
print(category_index)
{1: {'id': 1, 'name': 'player'}, 2: {'id': 2, 'name': 'ball'}}
def load_image_into_numpy_array(image):
(im_width, im_height) = image.size
return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
def best_rectangle(good_detections, image_width, image_height, scale = 0.5):
desired_width = int(image_width * scale)
desired_height = int(image_height * scale)
max_detections = 0
best_x = 0
best_y = 0
for i in range(0, int(image_width-desired_width)):
for j in range(0, int(image_height - desired_height)):
num_detections = 0
rect_x = (i, i + desired_width)
rect_y = (j , j+ desired_height)
for detection in good_detections:
if detection[0] >= rect_x[0] and detection[0] <= rect_x[1]:
if detection[1] >= rect_y[0] and detection[1] <= rect_y[1]:
# In this box
# print("In the box")
num_detections += 1
if num_detections > max_detections:
max_detections = num_detections
best_x = i
best_y = j
return max_detections, best_x, best_y
import math
def rotatePolygon(polygon,theta):
"""Rotates the given polygon which consists of corners represented as (x,y),
around the ORIGIN, clock-wise, theta degrees"""
theta = math.radians(theta)
rotatedPolygon = []
for corner in polygon :
rotatedPolygon.append(( corner[0]*math.cos(theta)-corner[1]*math.sin(theta) , corner[0]*math.sin(theta)+corner[1]*math.cos(theta)) )
return rotatedPolygon
from math import sin, cos, radians
def rotate_point(point, angle, center_point=(0, 0)):
"""Rotates a point around center_point(origin by default)
Angle is in degrees.
Rotation is counter-clockwise
"""
angle_rad = radians(angle % 360)
# Shift the point so that center_point becomes the origin
new_point = (point[0] - center_point[0], point[1] - center_point[1])
new_point = (new_point[0] * cos(angle_rad) - new_point[1] * sin(angle_rad),
new_point[0] * sin(angle_rad) + new_point[1] * cos(angle_rad))
# Reverse the shifting we have done
new_point = (new_point[0] + center_point[0], new_point[1] + center_point[1])
return new_point
def rotate_rectangle(coord1, coord2, good_detections, image_np):
angleRange = [-30, -15, 0, 15, 30]
max_detections = 0
best_x = 0
best_y = 0
best_angle = 0
tot_detections = []
for angle in angleRange:
print("On angle: ", angle)
new_point2 = rotate_point(coord2, angle, coord1)
newCoord1 = coord1
newCoord2 = (int(new_point2[0]), int(new_point2[1]))
print("Coordinates are: ", newCoord1, coord2, newCoord2)
#calculate num detections
num_detections = 0
for detection in good_detections:
if detection[0] >= newCoord1[0] and detection[0] <= newCoord2[0]:
if detection[1] >= newCoord1[1] and detection[1] <= newCoord2[1]:
# In this box
# print("In the box")
num_detections += 1
# for sanity checking
tot_detections.append(num_detections)
if num_detections > max_detections:
max_detections = num_detections
best_coord1 = newCoord1
best_coord2 = newCoord2
best_angle = angle
# cv2.rectangle(image_np, newCoord1, newCoord2 ,(0,0,255),3)
return max_detections, best_coord1, best_coord2, best_angle, tot_detections
filename = '/home/priya/Documents/AI_Apps/soccer_project/soccer_mini_slow.mp4'
out_filename = 'soccer_mini' + ".avi"
cap = cv2.VideoCapture(filename)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) # uses given video width and height
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
vwriter = cv2.VideoWriter(out_filename,cv2.VideoWriter_fourcc(*'MJPG'),fps, (1024, 600))
length = int(cap.get(7))
fps = cap.get(5)
duration = length/fps
# Running the tensorflow session
with detection_graph.as_default():
with tf.Session(graph=detection_graph) as sess:
counter = 0
fac = 3
while True:
ret, image_np = cap.read()
if not ret:
break
orig_image = image_np
origW, origH, _ = orig_image.shape
new_width = 1024
new_height = 600
image_np = cv2.resize(image_np, (new_width,new_height), interpolation=cv2.INTER_AREA)
counter += 1
if counter % fac == 0:
## Extract info from object detection
image_np_expanded = np.expand_dims(image_np, axis=0)
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
# Each box represents a part of the image where a particular object was detected.
boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represent how level of confidence for each of the objects.
# Score is shown on the result image, together with the class label.
scores = detection_graph.get_tensor_by_name('detection_scores:0')
classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name('num_detections:0')
# Actual detection.
(boxes, scores, classes, num_detections) = sess.run(
[boxes, scores, classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
# Visualization of the results of a detection.
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
np.squeeze(boxes),
np.squeeze(classes).astype(np.int32),
np.squeeze(scores),
category_index,
use_normalized_coordinates=True,
line_thickness=3,
min_score_thresh=0.6)
## Magic to draw box
good_detections = []
for n in range(len(scores[0])):
if scores[0][n] > 0.60:
# Calculate position
ymin = int(boxes[0][n][0] * new_height)
xmin = int(boxes[0][n][1] * new_width)
ymax = int(boxes[0][n][2] * new_height)
xmax = int(boxes[0][n][3] * new_width)
#crop them
x_avg = (xmin + xmax)/2
y_avg = (ymin + ymax)/2
if classes[0][n] == 1:
good_detections.append((x_avg, y_avg))
if classes[0][n] == 2:
print("Ball Detected")
good_detections.append((x_avg, y_avg))
good_detections.append((x_avg, y_avg))
good_detections.append((x_avg, y_avg))
good_detections.append((x_avg, y_avg))
good_detections.append((x_avg, y_avg))
desired_scale = 0.40
max_detections, best_x, best_y = best_rectangle(good_detections, new_width, new_height, desired_scale)
# print(max_detections, best_x, best_y)
coord1 = (int(best_x), int(best_y))
coord2 = (int(best_x + desired_scale*new_width), int(best_y + desired_scale*new_height))
# max_detections, best_coord1, best_coord2, best_angle, tot_detections = rotate_rectangle(coord1, coord2, good_detections, image_np)
# print("Detection by angle: ", tot_detections)
# print("Best angle is: ", best_angle)
cv2.rectangle(image_np, coord1, coord2 ,(255,0,0),2)
vwriter.write(image_np)
# cv2.imshow('image', image_np)
# key = cv2.waitKey(20) & 0xFF
# # if the `q` key is pressed, break from the lop
# if key == ord("q"):
# break
cap.release()
Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected Ball Detected
--------------------------------------------------------------------------- KeyboardInterrupt Traceback (most recent call last) <ipython-input-11-0c21095f00fb> in <module>() 85 86 desired_scale = 0.40 ---> 87 max_detections, best_x, best_y = best_rectangle(good_detections, new_width, new_height, desired_scale) 88 # print(max_detections, best_x, best_y) 89 coord1 = (int(best_x), int(best_y)) <ipython-input-8-141907888110> in best_rectangle(good_detections, image_width, image_height, scale) 10 num_detections = 0 11 rect_x = (i, i + desired_width) ---> 12 rect_y = (j , j+ desired_height) 13 for detection in good_detections: 14 if detection[0] >= rect_x[0] and detection[0] <= rect_x[1]: KeyboardInterrupt:
### Challenge is that boxes that are being drawn are parallel, not an angle.