from os import listdir
from os.path import isfile, join
from collections import deque
import math
import argparse
import numpy as np
import cv2
#определяем границы для красного и зеленого цвета
#красный цвет представляет из себя две области в пространстве HSV
red1 = np.array([0, 50, 20], dtype = "uint8")
red2 = np.array([5, 255, 255], dtype = "uint8")
red3 = np.array([175, 50, 20], dtype = "uint8")
red4 = np.array([180, 255, 255], dtype = "uint8")
#с зеленым все проще - он в центре диапазона
green1 = np.array([53, 90, 90], dtype = "uint8")
green2 = np.array([91, 255, 255], dtype = "uint8")
#путь по умолчанию к видеофайлам
path = 'C:\\Users\\dns\\Desktop\\tr\\'
#очищение файла с результатом, если он уже чем-то заполнен
f1 = open(path + 'res.txt', 'w')
f1.close()
#функция для вычисления расстояния
def distance(p1,p2):
return math.sqrt((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2 )
#функция, ищущая переключение
def findswitch(video_file):
cap = cv2.VideoCapture(video_file) #открытие видео
frame_len = 7
switches_array = deque(maxlen=frame_len)
true_switches_array = [] #массив переключений
offset = 10
x_mult_offset = 30
frame_num = 0 #номер текущего кадра
red_frames_arr = deque(maxlen=2*frame_len)
green_frames_arr = deque(maxlen=2*frame_len)
prev_frame = deque(maxlen=2*frame_len)
while(cap.isOpened()):
ret, frame = cap.read() #получаем изображение из видеопотока
f1 = open(path + 'res.txt', 'a')
if frame is None:
if len(true_switches_array) == 0:
print(filename, -1)
f1.write('%s %d\n' % (filename, -1))
break
switch_frame = frame_num #кадр переключения(текущий)
frame_num = frame_num + 1
frame_h,frame_w, _ = frame.shape #определяем высоту и ширину изображения
blurred = cv2.GaussianBlur(frame, (7, 7), 0.5) #сглаживание кадра
hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV) #изbgrвhsv
#красная и зелёная маски
red_mask = cv2.inRange(hsv, red1, red2) + cv2.inRange(hsv, red3, red4)
green_mask = cv2.inRange(hsv, green1, green2)
if len(prev_frame) < 1:
prev_frame.appendleft(blurred)
#определение разницы во времени между красным и зеленым сигналом
frameDelta = cv2.absdiff(prev_frame[len(prev_frame)-1], blurred)
prev_frame.appendleft(blurred)
diff_red_mask = np.zeros(frame_h*frame_w, dtype = "uint8").reshape(frame_h,frame_w)
if len(red_frames_arr) > 1 :
diff_red_mask = red_frames_arr[len(red_frames_arr) - 1] - red_mask
red_frames_arr.appendleft(red_mask)
diff_green_mask = np.zeros(frame_h*frame_w, dtype = "uint8").reshape(frame_h,frame_w)
if len(green_frames_arr) > 1 :
diff_green_mask = green_mask - green_frames_arr[len(green_frames_arr) - 1]
green_frames_arr.appendleft(green_mask)
diff_red_mask = cv2.erode(diff_red_mask, None, iterations=1)
diff_red_mask = cv2.dilate(diff_red_mask, None, iterations=3)
ret, diff_red_mask = cv2.threshold(diff_red_mask,127,250,cv2.THRESH_BINARY)
diff_green_mask = cv2.erode(diff_green_mask, None, iterations=1)
diff_green_mask = cv2.dilate(diff_green_mask, None, iterations=3)
ret, diff_green_mask = cv2.threshold(diff_green_mask,127,250,cv2.THRESH_BINARY)
red_contours = cv2.findContours(diff_red_mask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)[-2] #ищем контуры красных областей
green_contours = cv2.findContours(diff_green_mask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)[-2] #ищем контуры зеленых областей
area = 0
red_circle_array = [] #массив красных кругов
for red_contour in red_contours:
red_circle = cv2.minEnclosingCircle(red_contour)
((red_contour_x, red_contour_y), red_contour_radius) = red_circle
min_radius = 2 + (red_contour_x/frame_h) * 3
max_radius = 35 + (red_contour_x/frame_h) * 30
if red_contour_radius > min_radius and red_contour_radius < max_radius:
new_circle = (red_circle, area)
red_circle_array.append(new_circle) #добавление красного круга в массив, если он подходит по условиям
green_circle_array = [] #массив зеленых кругов
for green_contour in green_contours:
green_circle = cv2.minEnclosingCircle(green_contour)
((green_contour_x, green_contour_y), green_contour_radius) = green_circle
min_radius = 2 + (green_contour_x/frame_h) * 3
max_radius = 35 + (green_contour_y/frame_h) * 30
if green_contour_radius > min_radius and green_contour_radius < max_radius:
new_circle = (green_circle, area)
green_circle_array.append(new_circle) #добавление зелёного круга в массив, если он подходит по условиям
switches = [] #массив всех переключений
for red_circle in red_circle_array:
((red_x, red_y), red_contour_radius) = red_circle[0]
red_x = int(red_x)
red_y = int(red_y)
top = 0
left = 0
right = frameDelta.shape[1]
bottom = frameDelta.shape[0]
circle_offset = offset + red_contour_radius
if top < int(red_y - circle_offset):
top = int(red_y - circle_offset)
if left < int(red_x - (x_mult_offset*circle_offset)):
left = int(red_x - (x_mult_offset*circle_offset))
if bottom > int(red_y + circle_offset):
bottom = int(red_y + circle_offset)
if right > int(red_x + (x_mult_offset*circle_offset)):
right = int(red_x + (x_mult_offset*circle_offset))
red_crop = frameDelta[top:bottom,left:right,:]
total_move_green_red = np.sum(red_crop) - (380*red_contour_radius*red_contour_radius)
total_move = total_move_green_red / (red_crop.shape[0]*red_crop.shape[1])
#определение расстояний между кругами на кадре
for green_circle in green_circle_array:
dif_x = red_circle[0][0][0] - green_circle[0][0][0]
dif_y = red_circle[0][0][1] - green_circle[0][0][1]
dif_r = red_circle[0][1]/green_circle[0][1]
radius_red = red_circle[0][1]
max_dist_y = -5*(radius_red+green_circle[0][1])/2
if max_dist_y < -170:
max_dist_y = -170
if dif_r > 0.4 and dif_r < 2.5:
if dif_x < (radius_red/2) and dif_x > (-radius_red/2):
if dif_y > max_dist_y and dif_y < (-(radius_red+green_circle[0][1])/4) and dif_y < -7:
if total_move < 33:
new_switch = (red_circle[0], green_circle[0],total_move)
switches.append(new_switch)
#добавление верных переключений в нужный массив
for switch in switches:
true_switch = 0
frame_delta = 1
for last_switches in switches_array:
frame_delta = frame_delta + 1
for last_switch in last_switches:
distance_red = distance(switch[0][0],last_switch[0][0])
distance_green = distance(switch[1][0],last_switch[1][0])
if distance_red < (switch[0][1]*0.5) and distance_green < (switch[1][1]*0.5):
true_switch = true_switch + 1
if (frame_num - frame_delta) < switch_frame:
switch_frame = frame_num - frame_delta
break
if true_switch > 1:
true_switches_array.append(switch)
break
switches_array.appendleft(switches)
#отрисовка найденных красных и зелёных кругов на видео
for switch in switches:
cv2.circle(frame, (int(switch[0][0][0]), int(switch[0][0][1])), int(20), (255, 0, 255), 2)
for red_circle in red_circle_array:
cv2.circle(frame, (int(red_circle[0][0][0]), int(red_circle[0][0][1])), int(red_circle[0][1]), (0, 0, 255), 2)
for green_circle in green_circle_array:
cv2.circle(frame, (int(green_circle[0][0][0]), int(green_circle[0][0][1])), int(green_circle[0][1]), (0, 255, 0), 2)
#отрисовка кругов другого цвета при найденном нужном переключении
for switch in true_switches_array:
cv2.circle(frame, (int(switch[0][0][0]), int(switch[0][0][1])), int(switch[0][1]), (255, 0, 0), 2)
cv2.circle(frame, (int(switch[1][0][0]), int(switch[1][0][1])), int(switch[1][1]), (255, 255, 0), 2)
#вывод видео на экран
cv2.imshow("video", frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
#вывод в консоль и в файл
if len(true_switches_array) > 0:
print(filename, switch_frame)
f1.write('%s %05d\n' % (filename, switch_frame))
break
cap.release()
f1.close()
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("-f", "--folder", required=False, help="path to videos directory")
args = vars(ap.parse_args())
if args['folder'] is not None:
path = args['folder']
print(path)
#берем все файлы из пути с разрешением .avi и применяем к ним функцию поиска переключений
file_list = [f for f in listdir(path) if isfile(join(path, f)) and f.split('.')[-1] == 'avi']
for filename in file_list:
findswitch(path + filename)