mobilesam v1.4

app.py

@@ -15,9 +15,9 @@ from utils import (
 app = FastAPI()
 
 # Загрузка моделей
-sam_model = load_sam()       # MobileSAM
-model2 = load_model2()       # сорт
-model3 = load_model3()       # свежесть
+sam_predictor = load_sam()   # MobileSAM + SamPredictor
+model2 = load_model2()
+model3 = load_model3()
 
 DEVICE = torch.device('cpu')
 
@@ -39,13 +39,15 @@ async def predict_full(
     image = Image.open(io.BytesIO(content)).convert('RGB')
     orig_np = np.array(image)
 
-    # MobileSAM: сегментация по точке
+    # Установка изображения в MobileSAM Predictor
+    sam_predictor.set_image(orig_np)
+
+    # Промпт: точка на фрукте
     input_point = np.array([[point_x, point_y]])
     input_label = np.array([1])  # 1 = foreground
 
-    # Правильный вызов MobileSAM (из официального примера)
-    masks, scores, logits = sam_model.predict(
-        img=orig_np,
+    # Предсказание маски
+    masks, scores, logits = sam_predictor.predict(
         point_coords=input_point,
         point_labels=input_label,
         multimask_output=False

models.py

@@ -1,22 +1,23 @@
 import torch
 import torchvision.models as models
 import torch.nn as nn
-from mobile_sam import sam_model_registry
+from mobile_sam import sam_model_registry, SamPredictor
 
 DEVICE = torch.device('cpu')
 
-sam_model = None  # MobileSAM
-model2 = None     # сорт фрукта
-model3 = None     # свежесть
+sam_predictor = None  # MobileSAM + Predictor
+model2 = None         # сорт фрукта
+model3 = None         # свежесть
 
 def load_sam(weights_path='weights/mobile_sam.pt'):
-    global sam_model
-    if sam_model is None:
+    global sam_predictor
+    if sam_predictor is None:
         model_type = "vit_t"
-        sam_model = sam_model_registry[model_type](checkpoint=weights_path)
-        sam_model.to(DEVICE)
-        sam_model.eval()
-    return sam_model
+        sam = sam_model_registry[model_type](checkpoint=weights_path)
+        sam.to(DEVICE)
+        sam.eval()
+        sam_predictor = SamPredictor(sam)
+    return sam_predictor
 
 def load_model2(weights_path='weights/class.pth'):
     global model2

utils.py

@@ -17,7 +17,6 @@ def preprocess_for_classifier(img: np.ndarray) -> torch.Tensor:
     ])
     return transform(img)
 
-# Letterbox без искажения
 def letterbox_any_size(
     img: np.ndarray,
     target_size: int = 224,
@@ -40,10 +39,9 @@ def letterbox_any_size(
                                 cv2.BORDER_CONSTANT, value=bg_color)
     return padded
 
-# Обрезка по маске SAM + белый фон + letterbox
 def crop_fruit_with_white_bg(
-    orig_img: np.ndarray,      # RGB
-    mask: np.ndarray,          # bool от SAM
+    orig_img: np.ndarray,
+    mask: np.ndarray,
     out_size: int = 224,
     bg_color: tuple = (255, 255, 255)
 ) -> np.ndarray: