Update

app.py

@@ -50,12 +50,14 @@ def load_model(model_name: str):
     return feature_extractor, model
 
 
-def read_video(file_path: str) -> np.ndarray:
+def read_video(file_path: str, frames_per_video: int = 8) -> np.ndarray:
     cap = cv2.VideoCapture(file_path)
     length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 1000 frames
     print("Number of frames", length)
 
-    indices = sample_frame_indices(clip_len=8, frame_sample_rate=4, seg_len=length)
+    indices = sample_frame_indices(
+        clip_len=frames_per_video, frame_sample_rate=4, seg_len=length
+    )
 
     frames: List[np.array] = []
     for i in indices:
@@ -83,8 +85,8 @@ def read_video_decord(file_path: str) -> np.ndarray:
     return video
 
 
-def inference(file_path: str):
-    video = read_video(file_path)
+def inference(file_path: str, frames_per_video: int = 8):
+    video = read_video(file_path, frames_per_video)
 
     inputs = feature_extractor(list(video), return_tensors="pt")
 
@@ -111,6 +113,15 @@ def inference(file_path: str):
     return pd.DataFrame(results, columns=("Label", "Confidence"))
 
 
+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
+
+
 st.title("TimeSFormer")
 
 with st.expander("INTRODUCTION"):
@@ -135,6 +146,10 @@ model_name = st.selectbox(
 )
 feature_extractor, model = load_model(model_name)
 
+frames_per_video = get_frames_per_video(model_name)
+st.info(f"Frames per video: {frames_per_video}")
+
+
 VIDEO_TMP_PATH = os.path.join("tmp", "tmp.mp4")
 uploadedfile = st.file_uploader("Upload file", type=["mp4"])
 
@@ -146,23 +161,8 @@ if uploadedfile is not None:
 
     start_time = time.time()
     with st.spinner("Processing..."):
-        df = inference(VIDEO_TMP_PATH)
+        df = inference(VIDEO_TMP_PATH, frames_per_video)
     end_time = time.time()
     st.info(f"{end_time - start_time} seconds")
     st.dataframe(df)
     st.video(VIDEO_TMP_PATH)
-
-img_file_buffer = st.camera_input("Take a picture")
-
-if img_file_buffer is not None:
-    # To read image file buffer with OpenCV:
-    bytes_data = img_file_buffer.getvalue()
-    cv2_img = cv2.imdecode(np.frombuffer(bytes_data, np.uint8), cv2.IMREAD_COLOR)
-
-    # Check the type of cv2_img:
-    # Should output: <class 'numpy.ndarray'>
-    st.write(type(cv2_img))
-
-    # Check the shape of cv2_img:
-    # Should output shape: (height, width, channels)
-    st.write(cv2_img.shape)

capture_picture.py

@@ -0,0 +1,20 @@
+import streamlit as st
+import cv2
+import numpy as np
+
+img_file_buffer = st.camera_input("Take a picture")
+
+if img_file_buffer is not None:
+    # To read image file buffer with OpenCV:
+    bytes_data = img_file_buffer.getvalue()
+    cv2_img: np.ndarray = cv2.imdecode(
+        np.frombuffer(bytes_data, np.uint8), cv2.IMREAD_COLOR
+    )
+
+    # Check the type of cv2_img:
+    # Should output: <class 'numpy.ndarray'>
+    st.write(type(cv2_img))
+
+    # Check the shape of cv2_img:
+    # Should output shape: (height, width, channels)
+    st.write(cv2_img.shape)