init project

modules/pe3r/models.py

@@ -12,6 +12,7 @@ from modules.mobilesamv2 import sam_model_registry
 from sam2.sam2_video_predictor import SAM2VideoPredictor
 import spaces
 import torch
+from ultralytics import YOLOvv8
 
 class Models:
     @spaces.GPU
@@ -47,8 +48,9 @@ class Models:
         self.mobilesamv2.eval()
 
         # -- yolov8 --
-        YOLO8_CKP='./checkpoints/ObjectAwareModel.pt'
-        self.yolov8 = ObjectAwareModel(YOLO8_CKP)
+        # YOLO8_CKP='./checkpoints/ObjectAwareModel.pt'
+        # self.yolov8 = ObjectAwareModel(YOLO8_CKP)
+        self.yolov8 = YOLOvv8.from_pretrained("Ultralytics/YOLOv8")
 
         # -- siglip --
         self.siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)