lucky-verma · September 7, 2022 15:09 · Sep 7, 2022
diff --git a/templateMatchTrackEVAL.py b/templateMatchTrackEVAL.py
@@ -0,0 +1,127 @@
+import cv2
+import pathlib
+import time
+import argparse
+import os
+import glob
+
+# Construct the argument parser.
+parser = argparse.ArgumentParser()
+parser.add_argument('-i',
+                    '--input',
+                    help='path to input video',
+                    default='input/video_1.mp4')
+parser.add_argument('-t',
+                    '--template',
+                    help='path to the template',
+                    default='input/video_1_template.jpg')
+args = vars(parser.parse_args())
+
+print(args['template'].split(os.path.sep)[-1])
+
+res = []
+
+# Get all folders in the input folder.
+folder = args['template']
+
+sub_folders = [
+    name for name in os.listdir(folder)
+    if os.path.isdir(os.path.join(folder, name))
+]
+
+print(sub_folders)
+
+for f in sub_folders:
+
+    # Read the video input.
+    cap = cv2.VideoCapture(args['input'])
+    if (cap.isOpened() == False):
+        print('Error while trying to read video. Please check path again')
+    # Get the frame width and height
+    frame_width = int(cap.get(3))
+    frame_height = int(cap.get(4))
+    # String name with which to save the resulting video.
+    save_name = str(pathlib.Path(args['input'])).split(
+        os.path.sep)[-1].split('.')[0]
+    # define codec and create VideoWriter object
+    out = cv2.VideoWriter(f"outputs/{save_name}.mp4",
+                        cv2.VideoWriter_fourcc(*'mp4v'), 30,
+                        (frame_width, frame_height))
+
+    player_images = [[file, cv2.imread(file, 0)]
+                    for file in glob.glob('{}/*'.format(args['template'] + '/' + str(f)))]
+
+    for i in player_images:
+        print(i[0], i[0].split(os.path.sep)[-2].lstrip('0'))
+        # Read the template in grayscale format.
+        template = i[1]
+        w, h = template.shape[::-1]
+        frame_count = 0  # To count total frames.
+        total_fps = 0  # To get the final frames per second.
+
+        # get player id
+        # player_id = int(i[0].split(os.path.sep)[-2].lstrip('0'))
+        player_id = int(f.lstrip('0'))
+
+        # get the frame number from filename
+        number = int(i[0].split('_')[-2].lstrip("0"))
+        print(f'Player {player_id}', f'Frame {number}')
+
+        # Capture each frame of the video.
+        cap.set(1, number)
+        ret, frame = cap.read()
+        if ret:
+            image = frame.copy()
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            start_time = time.time()
+            # Apply template Matching.
+            result = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
+            end_time = time.time()
+
+            # Get the current fps.
+            fps = 1 / (end_time - start_time)
+            # Add `fps` to `total_fps`.
+            total_fps += fps
+            # Increment frame count.
+            frame_count += 1
+            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
+            # Top left x and y coordinates.
+            x1, y1 = max_loc
+            # Bottom right x and y coordinates.
+            x2, y2 = (x1 + w, y1 + h)
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
+            cv2.putText(frame,
+                        f"{fps:.1f} FPS", (15, 25),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        1, (0, 255, 0),
+                        2,
+                        lineType=cv2.LINE_AA)
+            cv2.imshow('Result', frame)
+            out.write(frame)
+
+
+            # TrackEval result.
+            temp = str(number) + ' ' + str(player_id) + ' ' + str(x1) + ' ' + str(y1) + ' ' + str(x2) + ' ' + str(y2) + ' -1, -1, -1, -1'
+            res.append(temp)
+            print(temp)
+
+            # Press `q` to exit.
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+        else:
+            break
+
+with open('{}.txt'.format(args['template']), 'w') as f:
+    for line in res:
+        f.write(f"{line}\n")
+
+
+# Release VideoCapture() object.
+cap.release()
+# Close all frames and video windows.
+cv2.destroyAllWindows()
+
+# python main.py --input input/video_1.mp4 --template input/video_1_template.jpg
+
+# Input: A directory containing the players & its template images.
+# python main.py --input "C:\Users\LuckyVerma\Documents\TemplateMatching\data\2000_plays\plays\run_right_seah_49s113.mp4" -t data\reid-v1\successful\labels-only\run_right_seah_49s113