129 lines
2.5 KiB
Python
129 lines
2.5 KiB
Python
import matplotlib.pyplot as plt
|
|
import saved_cordinats
|
|
|
|
|
|
from imageai.Detection import VideoObjectDetection
|
|
|
|
vid_obj_detect = VideoObjectDetection()
|
|
vid_obj_detect.setModelTypeAsYOLOv3()
|
|
|
|
vid_obj_detect.setModelPath(r"/home/mikolaj/2ait_tech/sportowe/film/yolo.h5")
|
|
vid_obj_detect.loadModel()
|
|
|
|
|
|
person_points = []
|
|
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'][3]))
|
|
else:
|
|
person_points.append((index, object['box_points'][0], object['box_points'][3]))
|
|
|
|
detected_vid_obj = vid_obj_detect.detectObjectsFromVideo(
|
|
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,
|
|
return_detected_frame = True,
|
|
)
|
|
|
|
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 = 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:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
1:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
2:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
3:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
4:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
5:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
6:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
7:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
8:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
9:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
10:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
11:{
|
|
'x': [],
|
|
'y': []
|
|
},
|
|
}
|
|
|
|
|
|
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():
|
|
|
|
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('shorted_wykres.png')
|
|
plt.show() |