film_sportowe/main.py

140 lines
2.8 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'][1]))
# else:
# person_points.append((index, object['box_points'][0], object['box_points'][1]))
# detected_vid_obj = vid_obj_detect.detectObjectsFromVideo(
# input_file_path = r"man.mp4",
# output_file_path = r"man_changed.mp4",
# frames_per_second=20,
# frame_detection_interval=2,
# 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)
ax.set_xlim([0, 960])
ax.set_ylim([540, 0])
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):
# if p[0] > 6:
# continue
if index - 1 >= 0:
if abs(p[1] - person_points[index -1][1]) > 20:
continue
if abs(p[2] - person_points[index -1][2]) > 20:
continue
persons[p[0]&7]['x'].append(p[1])
persons[p[0]&7]['y'].append(p[2])
for k, v in persons.items():
# if k > 6:
# continue
ax.plot(
persons[k]['x'],
persons[k]['y'],
color=COLORS[k%7],
lw=2,
label=LABELS[k%7]
)
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('man_shorted_wykres.png')
plt.show()