predict fix

.dockerignore

@@ -19,3 +19,25 @@ coverage.xml
 .mypy_cache
 .pytest_cache
 .hypothesis
+
+# generated by replicate/cog
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+.Python
+env
+pip-log.txt
+pip-delete-this-directory.txt
+.tox
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.log
+.git
+.mypy_cache
+.pytest_cache
+.hypothesis

predict.py

@@ -2,10 +2,13 @@ import os
 from cog import BasePredictor, Input, Path
 import torch
 import json
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
 from src.models.model import load_model
-from src.data.video_utils import create_transform, extract_frames
+from src.dataset.video_utils import create_transform, extract_frames
 
-CHECKPOINT_DIR = "runs/run_20241024-150232_otherpeopleval_large_model/"
+CHECKPOINT_DIR = "checkpoints/"
 
 class Predictor(BasePredictor):
     def setup(self):
@@ -24,7 +27,7 @@ class Predictor(BasePredictor):
         # Load model
         self.model = load_model(
             self.config['num_classes'],
-            os.path.join(CHECKPOINT_DIR, "best_model.pth"),
+            os.path.join(CHECKPOINT_DIR, "weights.ckpt"),
             self.device,
             self.config['clip_model']
         )