Download our e-book of Introduction To Python
Mahesh Pardeshi
3 years ago
class Model:
Class that contains the model and all its functions
"""
def __init__(self, model_path):
"""
Initialization function
@ model_path : path to the model
"""
# Declare detection graph
self.detection_graph = tf.Graph()
# Load the model into the tensorflow graph
with self.detection_graph.as_default():
od_graph_def =
tf.compat.v1.GraphDef()
with tf.io.gfile.GFile(model_path, 'rb') as file:
serialized_graph = file.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
# Create a session from the detection graph
self.sess = tf.compat.v1.Session(graph=self.detection_graph)
def predict(self,img):
"""
Get the predicition results on 1 frame
@ img : our img vector
"""
# Expand dimensions since the model expects images to have shape: [1, none, none, 3]
img_exp = np.expand_dims(img, axis=0)
# Pass the inputs and outputs to the session to get the results
(boxes, scores, classes) = self.sess.run([self.detection_graph.get_tensor_by_name('detection_boxes:0'),
self.detection_graph.get_tensor_by_name('detection_scores:0'), self.detection_graph.get_tensor_by_name('detection_classes:0')],feed_dict={self.detection_graph.get_tensor_by_name('image_tensor:0'):
img_exp})
return (boxes, scores, classes)
def compute_perspective_transform(corner_points,width,height,image):
"""
Compute the transformation matrix
@corner_points : 4 corner points selected from the image
@height, width : size of the image
return : transformation matrix and the transformed image
"""
#Create an array out of the 4 corner points
corner_points_array = np.float32(corner_points)
#Create an array with the parameters (the dimensions) required to build the matrix
img_params = np.float32([[0,0],[width,0],[0,height],[width,height]])
# Compute and return the transformation matrix
matrix = cv2.getPerspectiveTransform(corner_points_array,img_params)
img_transformed = cv2.warpPerspective(image,matrix,(width,height))
return matrix,img_transformed
def compute_point_perspective_transformation(matrix,list_downoids):
"""
Apply the perspective transformation to every ground point which have been
detected on the main frame.
@matrix : the 3x3 matrix
@list_downoids : list that contains the points to transform
return: list containing all the new points
"""
#Compute the new coordinates of our points
list_points_to_detect= np.float32(list_downoids).reshape(-1, 1, 2)
transformed_points= cv2.perspectiveTransform(list_points_to_detect, matrix)
#Loop over the points and add them to the list that will be returned
transformed_points_list= list()
for i in range(0,transformed_points.shape[0]):
transformed_points_list.append([transformed_points[i][0][0],transformed_points[i][0][1]])
return transformed_points_list
# Check if 2 or more people have been detected (otherwise no need to detect)
if len(transformed_downoids) >= 2:
# Iterate over every possible 2 by 2 between the points combinations
list_indexes = list(itertools.combinations(range(len(transformed_downoids)), 2))
for i,pair in enumerate(itertools.combinations(transformed_downoids, r=2)):
# Check if the distance between each combination of points is less than the minimum distance chosen
if math.sqrt( (pair[0][0] - pair[1][0])**2 + (pair[0][1] - pair[1][1])**2 ) < int(distance_minimum):
# Change the colors of the points that are too close from each other to red
change_color_topview(pair)
# Get the equivalent indexes of these points in the original frame and change the color to red
index_pt1 = list_indexes[i][0]
index_pt2 = list_indexes[i][1]
change_color_originalframe(index_pt1,index_pt2)