vitpose-small

app.py

@@ -36,8 +36,8 @@ def process_video(video_file, threshold_secs):
     # Load ViTPose components from Hugging Face
     # This model is used to refine detection by analyzing keypoints in the cropped person area.
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    image_processor = AutoImageProcessor.from_pretrained("usyd-community/vitpose-base-simple")
-    vitpose_model = ViTPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-base-simple").to(device)
+    image_processor = AutoImageProcessor.from_pretrained("usyd-community/vitpose-plus-small")
+    vitpose_model = ViTPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-plus-small").to(device)
     vitpose_model.eval()
     
     # Dictionary to keep track of when a given track was first detected as "lying"