diff --git a/main.py b/main.py index f08c398..43b8f55 100644 --- a/main.py +++ b/main.py @@ -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]) + +for k, v in persons.items(): - 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 + 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() \ No newline at end of file