modified treatment script

src/social_detection_video_detection.py

@@ -1,12 +1,11 @@
-from bird_view_transfo import compute_perspective_transform
+from bird_view_transfo import compute_perspective_transform,compute_point_perspective_transformation
 from adrian_detection import detect_people
 import numpy as np
 import itertools
 import imutils
 import math
 import yaml
 import cv2
-import os
 
 
 """ Load the config parameters from the YAML file
@@ -25,7 +24,6 @@
 	height_og = int(cfg["image_parameters"]["height_og"])
 	img_path = cfg["image_parameters"]["img_path"]
 
-
 """ Load the YOLO weights and the config parameter
 """
 print("[ Loading YOLO model ] ...")
@@ -48,10 +46,12 @@
 # 				START THE VIDEO STREAM               #
 #########									 #########
 ######################################################
-vs = cv2.VideoCapture("../video/pedestrians.mp4")
+video_name = input("Enter the exact name of the video (including .mp4 or else) : ")
+
+vs = cv2.VideoCapture("../video/"+video_name)
+output_video_1,output_video_2 = None,None
 # Loop until the end of the video stream
 while True:
-
 	# Create a full black frame 
 	blank_image = np.zeros((height,width,3), np.uint8)
 
@@ -60,6 +60,7 @@
 	if not frame_exists:
 		break
 	else:
+		frame = imutils.resize(frame, width=1000)
 		# Detect the person in the frame and test if there is more 
 		results = detect_people(frame, net, ln, 0)
 
@@ -72,39 +73,53 @@
 				# initialize the color of the annotation
 				(startX, startY, endX, endY) = bbox
 				(cX, cY) = centroid
-				color = (0, 255, 0)
 				dist_x = int(math.sqrt((startX - endX)**2)/2)
 				dist_y = int(math.sqrt((startY - endY)**2)/2)
-
-				#cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
-				# cv2.circle(frame, (cX, cY+dist_y), 5, color, 1)
+				cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2)
+				# cv2.circle(frame, (cX, cY+dist_y), 5, (0, 255, 0), 1)
 				list_downoids.append([cX, cY+dist_y])
 
+			transformed_downoids = compute_point_perspective_transformation(matrix,list_downoids)
+			for point in transformed_downoids:
+				x,y = point
+				cv2.circle(blank_image, (x,y), 20, (0, 255, 0), 2)
+				cv2.circle(blank_image, (x,y), 3, (0, 255, 0), -1)
 
-			list_indexes = list(itertools.combinations(range(len(list_downoids)), 2))
-			for i,pair in enumerate(itertools.combinations(list_downoids, r=2)):
-				# print("----")
-				# print(pair)
-				# print(math.sqrt((pair[0][0]-pair[1][0])**2+(pair[0][1]-pair[1][1])**2))
-				if math.sqrt( (pair[0][0] - pair[1][0])**2 + (pair[0][1] - pair[1][1])**2 ) < 100:
+
+			list_indexes = list(itertools.combinations(range(len(transformed_downoids)), 2))
+			for i,pair in enumerate(itertools.combinations(transformed_downoids, r=2)):
+				if math.sqrt( (pair[0][0] - pair[1][0])**2 + (pair[0][1] - pair[1][1])**2 ) < 50:
+					cv2.circle(blank_image, (pair[0][0],pair[0][1]), 20, (0, 0, 255), 2)
+					cv2.circle(blank_image, (pair[0][0],pair[0][1]), 3, (0, 0, 255), -1)
+					cv2.circle(blank_image, (pair[1][0],pair[1][1]), 20, (0, 0, 255), 2)
+					cv2.circle(blank_image, (pair[1][0],pair[1][1]), 3, (0, 0, 255), -1)
 					index_pt1 = list_indexes[i][0]
 					index_pt2 = list_indexes[i][1]
 					(startX, startY, endX, endY) = results[index_pt1][1]
-					cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
+					cv2.rectangle(frame, (startX, startY), (endX, endY), ( 0, 0,255), 2)
 					(startX, startY, endX, endY) = results[index_pt2][1]
-					cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
+					cv2.rectangle(frame, (startX, startY), (endX, endY), ( 0, 0,255), 2)
 					(x1, y1) = results[index_pt1][2]
 					(x2, y2) = results[index_pt2][2]
 
 
 
 	cv2.imshow("Bird view", blank_image)
-	cv2.imwrite("blank_image.jpg", blank_image)
 	cv2.imshow("Frame", frame)
-	cv2.imwrite("Frame.jpg", frame)
 	key = cv2.waitKey(1) & 0xFF
 
+	if output_video_1 is None:
+		fourcc1 = cv2.VideoWriter_fourcc(*"MJPG")
+		output_video_1 = cv2.VideoWriter("output_1.avi", fourcc1, 25,(frame.shape[1], frame.shape[0]), True)
+	elif output_video_1 is not None:
+		output_video_1.write(frame)
+
+	if output_video_2 is None:	
+		fourcc2 = cv2.VideoWriter_fourcc(*"MJPG")
+		output_video_2 = cv2.VideoWriter("output_2.avi", fourcc2, 25,(blank_image.shape[1], blank_image.shape[0]), True)
+	elif output_video_2 is not None:
+		output_video_2.write(blank_image)
+
 	# if the `q` key was pressed, break from the loop
 	if key == ord("q"):
-		break
-	#break
+		break