Refactor

requirements.txt

@@ -6,3 +6,4 @@ black
 opencv-python
 opencv-python-headless
 streamlit-webrtc
+typed-argument-parser

run_opencv.py

@@ -1,134 +1,148 @@
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import cv2
 import numpy as np
 import pandas as pd
 import torch
+from tap import Tap
 from torch import Tensor
 from transformers import AutoFeatureExtractor, TimesformerForVideoClassification
 
 from utils.img_container import ImgContainer
 
 
-def load_model(model_name: str):
-    if "base-finetuned-k400" in model_name or "base-finetuned-k600" in model_name:
-        feature_extractor = AutoFeatureExtractor.from_pretrained(
-            "MCG-NJU/videomae-base-finetuned-kinetics"
-        )
-    else:
-        feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-    model = TimesformerForVideoClassification.from_pretrained(model_name)
-    return feature_extractor, model
+class ArgParser(Tap):
+    is_recording: Optional[bool] = False
+
+    # "facebook/timesformer-base-finetuned-k400"
+    # "facebook/timesformer-base-finetuned-k600",
+    # "facebook/timesformer-base-finetuned-ssv2",
+    # "facebook/timesformer-hr-finetuned-k600",
+    # "facebook/timesformer-hr-finetuned-k400",
+    # "facebook/timesformer-hr-finetuned-ssv2",
+    # "fcakyon/timesformer-large-finetuned-k400",
+    # "fcakyon/timesformer-large-finetuned-k600",
+    model_name: Optional[str] = "facebook/timesformer-base-finetuned-k400"
 
+    num_skip_frames: Optional[int] = 4
 
-def inference():
-    if not img_container.ready:
-        return
 
-    inputs = feature_extractor(list(img_container.imgs), return_tensors="pt")
+class ActivityModel:
+    def __init__(self, args: ArgParser):
+        self.feature_extractor, self.model = self.load_model(args.model_name)
 
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits: Tensor = outputs.logits
+        self.frames_per_video = self.get_frames_per_video(args.model_name)
+        print(f"Frames per video: {self.frames_per_video}")
 
-    # model predicts one of the 400 Kinetics-400 classes
-    max_index = logits.argmax(-1).item()
-    predicted_label = model.config.id2label[max_index]
+    def load_model(self, model_name: str):
+        if "base-finetuned-k400" in model_name or "base-finetuned-k600" in model_name:
+            feature_extractor = AutoFeatureExtractor.from_pretrained(
+                "MCG-NJU/videomae-base-finetuned-kinetics"
+            )
+        else:
+            feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+        model = TimesformerForVideoClassification.from_pretrained(model_name)
+        return feature_extractor, model
 
-    img_container.frame_rate.label = f"{predicted_label}_{logits[0][max_index]:.2f}%"
+    def inference(self, img_container: ImgContainer):
+        if not img_container.ready:
+            return
 
-    TOP_K = 12
-    # logits = np.squeeze(logits)
-    logits = logits.squeeze().numpy()
-    indices = np.argsort(logits)[::-1][:TOP_K]
-    values = logits[indices]
+        inputs = self.feature_extractor(list(img_container.imgs), return_tensors="pt")
 
-    results: List[Tuple[str, float]] = []
-    for index, value in zip(indices, values):
-        predicted_label = model.config.id2label[index]
-        # print(f"Label: {predicted_label} - {value:.2f}%")
-        results.append((predicted_label, value))
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits: Tensor = outputs.logits
 
-    img_container.rs = pd.DataFrame(results, columns=("Label", "Confidence"))
+        # model predicts one of the 400 Kinetics-400 classes
+        max_index = logits.argmax(-1).item()
+        predicted_label = self.model.config.id2label[max_index]
+
+        img_container.frame_rate.label = (
+            f"{predicted_label}_{logits[0][max_index]:.2f}%"
+        )
 
+        TOP_K = 12
+        # logits = np.squeeze(logits)
+        logits = logits.squeeze().numpy()
+        indices = np.argsort(logits)[::-1][:TOP_K]
+        values = logits[indices]
 
-def get_frames_per_video(model_name: str) -> int:
-    if "base-finetuned" in model_name:
-        return 8
-    elif "hr-finetuned" in model_name:
-        return 16
-    else:
-        return 96
+        results: List[Tuple[str, float]] = []
+        for index, value in zip(indices, values):
+            predicted_label = self.model.config.id2label[index]
+            # print(f"Label: {predicted_label} - {value:.2f}%")
+            results.append((predicted_label, value))
 
+        img_container.rs = pd.DataFrame(results, columns=("Label", "Confidence"))
 
-model_name = "facebook/timesformer-base-finetuned-k400"
-# "facebook/timesformer-base-finetuned-k400"
-# "facebook/timesformer-base-finetuned-k600",
-# "facebook/timesformer-base-finetuned-ssv2",
-# "facebook/timesformer-hr-finetuned-k600",
-# "facebook/timesformer-hr-finetuned-k400",
-# "facebook/timesformer-hr-finetuned-ssv2",
-# "fcakyon/timesformer-large-finetuned-k400",
-# "fcakyon/timesformer-large-finetuned-k600",
-feature_extractor, model = load_model(model_name)
+    def get_frames_per_video(self, model_name: str) -> int:
+        if "base-finetuned" in model_name:
+            return 8
+        elif "hr-finetuned" in model_name:
+            return 16
+        else:
+            return 96
 
 
-frames_per_video = get_frames_per_video(model_name)
-print(f"Frames per video: {frames_per_video}")
+def main(args: ArgParser):
+    activity_model = ActivityModel(args)
+    img_container = ImgContainer(activity_model.frames_per_video)
 
-img_container = ImgContainer(frames_per_video)
+    num_skips = 0
 
-SKIP_FRAMES = 4
+    # define a video capture object
+    camera = cv2.VideoCapture(0)
 
-num_skips = 0
+    frame_width = int(camera.get(3))
+    frame_height = int(camera.get(4))
+    size = (frame_width, frame_height)
 
-# define a video capture object
-camera = cv2.VideoCapture(0)
+    video_output = cv2.VideoWriter(
+        "activities.mp4", cv2.VideoWriter_fourcc(*"MP4V"), 10, size
+    )
 
-frame_width = int(camera.get(3))
-frame_height = int(camera.get(4))
-size = (frame_width, frame_height)
+    if camera.isOpened() == False:
+        print("Error reading video file")
 
-video_output = cv2.VideoWriter(
-    "activities.mp4", cv2.VideoWriter_fourcc(*"MJPG"), 10, size
-)
+    while camera.isOpened():
+        # Capture the video frame
+        # by frame
+        ret, frame = camera.read()
 
-if camera.isOpened() == False:
-    print("Error reading video file")
+        num_skips = (num_skips + 1) % args.num_skip_frames
 
-while camera.isOpened():
-    # Capture the video frame
-    # by frame
-    ret, frame = camera.read()
+        img_container.img = frame
+        img_container.frame_rate.count()
 
-    num_skips = (num_skips + 1) % SKIP_FRAMES
+        if num_skips == 0:
+            img_container.add_frame(frame)
+            activity_model.inference(img_container)
+        rs = img_container.frame_rate.show_fps(frame, img_container.is_recording)
 
-    img_container.img = frame
-    img_container.frame_rate.count()
+        # Display the resulting frame
+        cv2.imshow("ActivityTracking", rs)
 
-    if num_skips == 0:
-        img_container.add_frame(frame)
-        # inference()
-    rs = img_container.frame_rate.show_fps(frame, img_container.is_recording)
+        if img_container.is_recording:
+            video_output.write(rs)
 
-    # Display the resulting frame
-    cv2.imshow("ActivityTracking", rs)
+        # the 'q' button is set as the
+        # quitting button you may use any
+        # desired button of your choice
+        k = cv2.waitKey(1)
 
-    if img_container.is_recording:
-        video_output.write(rs)
+        if k == ord("q"):
+            break
+        elif k == ord("r"):
+            img_container.toggle_recording()
 
-    # the 'q' button is set as the
-    # quitting button you may use any
-    # desired button of your choice
-    k = cv2.waitKey(1)
+    # After the loop release the cap object
+    camera.release()
+    video_output.release()
+    # Destroy all the windows
+    cv2.destroyAllWindows()
 
-    if k == ord("q"):
-        break
-    elif k == ord("r"):
-        img_container.toggle_recording()
 
-# After the loop release the cap object
-camera.release()
-video_output.release()
-# Destroy all the windows
-cv2.destroyAllWindows()
+if __name__ == "__main__":
+    args = ArgParser().parse_args()
+    main(args)