Refactor

run_opencv.py

@@ -1,13 +1,14 @@
-from typing import List, Optional, Tuple
+from typing import List, Tuple
+
 import cv2
-from pandas import DataFrame
-from transformers import AutoFeatureExtractor, TimesformerForVideoClassification
 import numpy as np
-import torch 
 import pandas as pd
+import torch
 from torch import Tensor
+from transformers import AutoFeatureExtractor, TimesformerForVideoClassification
+
+from utils.img_container import ImgContainer
 
-from utils.frame_rate import FrameRate
 
 def load_model(model_name: str):
     if "base-finetuned-k400" in model_name or "base-finetuned-k600" in model_name:
@@ -19,23 +20,6 @@ def load_model(model_name: str):
     model = TimesformerForVideoClassification.from_pretrained(model_name)
     return feature_extractor, model
 
-class ImgContainer:
-    def __init__(self, frames_per_video: int = 8) -> None:
-        self.img: Optional[np.ndarray] = None  # raw image
-        self.frame_rate: FrameRate = FrameRate()
-        self.imgs: List[np.ndarray] = []
-        self.frame_rate.reset()
-        self.frames_per_video = frames_per_video
-        self.rs: Optional[DataFrame] = None
-
-    def add_frame(self, frame: np.ndarray):
-        if len(img_container.imgs) >= frames_per_video:
-            self.imgs.pop(0)
-        self.imgs.append(frame)
-
-    @property
-    def ready(self):
-        return len(img_container.imgs) == self.frames_per_video
 
 def inference():
     if not img_container.ready:
@@ -50,7 +34,7 @@ def inference():
     # model predicts one of the 400 Kinetics-400 classes
     max_index = logits.argmax(-1).item()
     predicted_label = model.config.id2label[max_index]
-    
+
     img_container.frame_rate.label = f"{predicted_label}_{logits[0][max_index]:.2f}%"
 
     TOP_K = 12
@@ -67,6 +51,7 @@ def inference():
 
     img_container.rs = pd.DataFrame(results, columns=("Label", "Confidence"))
 
+
 def get_frames_per_video(model_name: str) -> int:
     if "base-finetuned" in model_name:
         return 8
@@ -100,7 +85,7 @@ num_skips = 0
 # define a video capture object
 vid = cv2.VideoCapture(0)
 
-while(True):
+while True:
     # Capture the video frame
     # by frame
     ret, frame = vid.read()
@@ -109,19 +94,19 @@ while(True):
 
     img_container.img = frame
     img_container.frame_rate.count()
-    
+
     if num_skips == 0:
         img_container.add_frame(frame)
         inference()
     rs = img_container.frame_rate.show_fps(frame)
 
     # Display the resulting frame
-    cv2.imshow('TimeSFormer', rs)
+    cv2.imshow("TimeSFormer", rs)
 
     # the 'q' button is set as the
     # quitting button you may use any
     # desired button of your choice
-    if cv2.waitKey(1) & 0xFF == ord('q'):
+    if cv2.waitKey(1) & 0xFF == ord("q"):
         break
 
 # After the loop release the cap object

utils/frame_rate.py

@@ -1,6 +1,8 @@
+import time
 from typing import Optional
+
+import cv2
 import numpy as np
-import time, cv2
 
 
 class FrameRate:

utils/img_container.py

@@ -0,0 +1,24 @@
+from typing import List, Optional
+
+import numpy as np
+from pandas import DataFrame
+
+from .frame_rate import FrameRate
+
+
+class ImgContainer:
+    def __init__(self, frames_per_video: int = 8) -> None:
+        self.img: Optional[np.ndarray] = None  # raw image
+        self.frame_rate: FrameRate = FrameRate()
+        self.imgs: List[np.ndarray] = []
+        self.frames_per_video = frames_per_video
+        self.rs: Optional[DataFrame] = None
+
+    def add_frame(self, frame: np.ndarray) -> None:
+        if len(self.imgs) >= self.frames_per_video:
+            self.imgs.pop(0)
+        self.imgs.append(frame)
+
+    @property
+    def ready(self):
+        return len(self.imgs) == self.frames_per_video