commit 07efc9081e7b08601f41b0e8290694705ff6a699 Author: MikoĊ‚aj Pokrywka Date: Sun Jan 29 14:29:36 2023 +0100 wip diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b6cd6a4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +yolo.h5 +venv \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..27d7495 --- /dev/null +++ b/README.md @@ -0,0 +1 @@ +download yolo model https://github.com/OlafenwaMoses/ImageAI/releases/tag/1.0/ \ No newline at end of file diff --git a/changed_football.mp4.avi b/changed_football.mp4.avi new file mode 100644 index 0000000..6969589 Binary files /dev/null and b/changed_football.mp4.avi differ diff --git a/football.mp4 b/football.mp4 new file mode 100644 index 0000000..bc368db Binary files /dev/null and b/football.mp4 differ diff --git a/main.py b/main.py new file mode 100644 index 0000000..f08c398 --- /dev/null +++ b/main.py @@ -0,0 +1,94 @@ +from imageai.Detection import VideoObjectDetection +import matplotlib.pyplot as plt + + +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"football.mp4", + output_file_path = r"changed_football.mp4", + frames_per_second=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(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'] + +persons = { + 0:{ + 'x': [], + 'y': [] + }, + 1:{ + 'x': [], + 'y': [] + }, + 2:{ + 'x': [], + 'y': [] + }, + 3:{ + 'x': [], + 'y': [] + }, + 4:{ + 'x': [], + 'y': [] + }, + 5:{ + 'x': [], + 'y': [] + }, + # 6:{ + # '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]) + + 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 + +ax.legend() + +plt.savefig('wykres.png') +plt.show() \ No newline at end of file diff --git a/req.txt b/req.txt new file mode 100644 index 0000000..f7d09b5 --- /dev/null +++ b/req.txt @@ -0,0 +1,6 @@ +imageAI==2.1.6 +opencv-python +Pillow +numpy==1.21 +tensorflow +torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu