-
Notifications
You must be signed in to change notification settings - Fork 0
/
video_detector.py
81 lines (56 loc) · 2.18 KB
/
video_detector.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
import numpy as np
import cv2
import torch
import pandas as pd
import sys
TEXT_FONT = cv2.FONT_HERSHEY_PLAIN
TEXT_SCALE = 1.2
TEXT_THICKNESS = 1
TEXT_COLOR = (0, 255, 0)
DEBUG = False
#inizializzazione video
video_path = sys.argv[2]
stream = cv2.VideoCapture(video_path)
prev_frame_time = 0
new_frame_time = 0
weight_path = sys.argv[1]
model = torch.hub.load('ultralytics/yolov5', 'custom', path=weight_path)
model.eval()
# Idea: mandare in input alla rete l'immagine filtrata (passa-basso / passa-banda)
while True:
(sourceAvailable, frame) = stream.read()
if not sourceAvailable:
print("Source video non disponibile.")
break
#frame = cv2.flip(frame, 1)
result = model(frame)
pandas = result.pandas()
grabbed = False #non hai afferrato nessun oggetto
if len(pandas.xyxy) > 0:
for res in pandas.xyxy: #per ogni bounding box
for i in range(len(res)):
pt1 = (int(res["xmin"][i]), int(res["ymin"][i]))
pt2 = (int(res["xmax"][i]), int(res["ymax"][i]))
className = res["name"][i]
coloreBox = (255, 0, 0)
coloreBaricentro = (0, 0, 255)
if(className == "Nine_VFR" or className == "Nine_VFL"):
coloreBaricentro = (0, 255, 0)
coloreBox = (0, 255, 0)
grabbed = True
else:
grabbed = False
baricentroX = (pt2[0] + pt1[0]) // 2
baricentroY = (pt1[1] + pt2[1]) // 2
baricentro = (baricentroX, baricentroY)
cv2.rectangle(frame, pt1, pt2, color=coloreBox, thickness=2)
cv2.circle(frame, baricentro, 1, color=coloreBaricentro, thickness=4)
if DEBUG:
font_size = cv2.getTextSize(className, TEXT_FONT, TEXT_SCALE, TEXT_THICKNESS)
textP1 = pt1[0] + 5
textP2 = pt1[1] + font_size[1] - 10
cv2.putText(frame, className, (textP1, textP2), TEXT_FONT, TEXT_SCALE, TEXT_COLOR, TEXT_THICKNESS)
cv2.imshow("Frame", frame)
if(cv2.waitKey(1) & 0xFF == ord('q')):
break
stream.release()