This commit is contained in:
Mikołaj Pokrywka 2023-01-30 12:27:42 +01:00
parent 07efc9081e
commit 1a7e485599

85
main.py
View File

@ -1,7 +1,9 @@
from imageai.Detection import VideoObjectDetection
import matplotlib.pyplot as plt
import saved_cordinats
from imageai.Detection import VideoObjectDetection
vid_obj_detect = VideoObjectDetection()
vid_obj_detect.setModelTypeAsYOLOv3()
@ -16,14 +18,15 @@ ball_points = []
def forFrame(frame_number, output_array, output_count, detected_frame):
for index, object in enumerate(output_array):
if object['name'] == 'sports ball':
ball_points.append((index, object['box_points'][0], object['box_points'][1]))
ball_points.append((index, object['box_points'][0], object['box_points'][3]))
else:
person_points.append((index, object['box_points'][0], object['box_points'][1]))
person_points.append((index, object['box_points'][0], object['box_points'][3]))
detected_vid_obj = vid_obj_detect.detectObjectsFromVideo(
input_file_path = r"football.mp4",
output_file_path = r"changed_football.mp4",
input_file_path = r"shorted.mp4",
output_file_path = r"changed_shorted.mp4",
frames_per_second=20,
frame_detection_interval=20,
log_progress=True,
per_frame_function = forFrame,
# per_second_function=forSeconds,
@ -32,16 +35,18 @@ detected_vid_obj = vid_obj_detect.detectObjectsFromVideo(
with open('saved_cordinats', 'w') as file:
file.write(str(person_points))
file.write("\n"*3)
file.write(str(ball_points))
fig = plt.figure()
ax = fig.add_subplot(111)
# person_points = [(0, 206, 376), (1, 1166, 405), (2, 42, 362), (3, 710, 390), (4, 754, 408), (5, 1083, 378), (0, 216, 383), (1, 64, 367), (2, 709, 387), (3, 750, 400), (4, 1087, 374), (5, 1179, 389)]
# ball_points= [(6, 204, 564), (6, 211, 566), (6, 225, 563)]
COLORS = ['red', 'blue', 'orange', 'pink', 'yellow', 'brown']
LABELS = ['p1', 'p2', 'p3', 'p4', 'p5', 'p6']
# person_points = saved_cordinats.persons
# ball_points= saved_cordinats.ball
COLORS = ['red', 'blue', 'orange', 'pink', 'yellow', 'brown', 'white', 'magenta', 'cyan']
LABELS = ['p1', 'p2', 'p3', 'p4', 'p5', 'p6', 'p7', 'p8', 'p9', 'p10']
persons = {
0:{
@ -68,27 +73,57 @@ persons = {
'x': [],
'y': []
},
# 6:{
# 'x': [],
# 'y': []
# },
6:{
'x': [],
'y': []
},
7:{
'x': [],
'y': []
},
8:{
'x': [],
'y': []
},
9:{
'x': [],
'y': []
},
10:{
'x': [],
'y': []
},
11:{
'x': [],
'y': []
},
}
try:
for index, p in enumerate(person_points):
persons[p[0]]['x'].append(p[1])
persons[p[0]]['y'].append(p[2])
for k, v in persons.items():
ax.plot( persons[k]['x'], persons[k]['y'], color=COLORS[k], lw=2, label=LABELS[k])
for index, p in enumerate(person_points):
persons[p[0]]['x'].append(p[1])
persons[p[0]]['y'].append(p[2])
ball_x = [x[1] for x in ball_points]
ball_y = [x[2] for x in ball_points]
ax.plot(ball_x, ball_y, color='green', lw=2, label='ball')
except:
pass
for k, v in persons.items():
if k > 2:
continue
ax.plot(
persons[k]['x'],
persons[k]['y'],
color=COLORS[k],
lw=2,
label=LABELS[k]
)
ball_x = [x[1] for x in ball_points]
ball_y = [x[2] for x in ball_points]
ax.plot(ball_x, ball_y, color='green', lw=2, label='ball')
ax.legend()
plt.savefig('wykres.png')
plt.savefig('shorted_wykres.png')
plt.show()