Initialisation 7

app.py

@@ -4,7 +4,31 @@ import os
 import shutil
 import numpy as np
 from PIL import Image
-from utils.predictor import Predictor
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+class Predictor:
+    def __init__(self, model_cfg, checkpoint, device):
+        self.device = device
+        self.model = build_sam2(model_cfg, checkpoint, device=device)
+        self.predictor = SAM2ImagePredictor(self.model)
+        self.image_set = False
+
+    def set_image(self, image):
+        """Set the image for SAM prediction."""
+        self.image = image
+        self.predictor.set_image(image)
+        self.image_set = True
+
+    def predict(self, point_coords, point_labels, multimask_output=False):
+        """Run SAM prediction."""
+        if not self.image_set:
+            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
+        return self.predictor.predict(
+            point_coords=point_coords,
+            point_labels=point_labels,
+            multimask_output=multimask_output
+        )
 from utils.helpers import (
     blend_mask_with_image,
     save_mask_as_png,