mobilesam v1.1

app.py

@@ -2,9 +2,9 @@ from fastapi import FastAPI, UploadFile, File, Query
 import torch
 import numpy as np
 from PIL import Image
-import base64
 import io
-from models import load_model1, load_model2, load_model3
+import base64
+from models import load_sam, load_model2, load_model3
 from utils import (
     crop_fruit_with_white_bg,
     preprocess_for_classifier,
@@ -15,9 +15,9 @@ from utils import (
 app = FastAPI()
 
 # Загрузка моделей
-sam_predictor = load_model1()  # MobileSAM
-model2 = load_model2()
-model3 = load_model3()
+sam_model = load_sam()       # MobileSAM
+model2 = load_model2()       # сорт
+model3 = load_model3()       # свежесть
 
 DEVICE = torch.device('cpu')
 
@@ -32,23 +32,22 @@ async def predict_full(
     file: UploadFile = File(...),
     point_x: int = Query(..., description="X-координата точки на фрукте (в пикселях оригинального изображения)"),
     point_y: int = Query(..., description="Y-координата точки на фрукте"),
-    return_cropped: bool = Query(default=True, description="Вернуть обрезанное изображение в base64?")
+    return_cropped: bool = Query(default=True, description="Вернуть обрезанное изображение в base64?"),
+    cropped_size: int = Query(224, description="Размер обрезанного изображения (100 или 224)")
 ):
     content = await file.read()
     image = Image.open(io.BytesIO(content)).convert('RGB')
     orig_np = np.array(image)
 
-    # Установка изображения в SAM
-    sam_predictor.set_image(orig_np)
-
-    # Промпт: точка на фрукте
+    # MobileSAM: сегментация по точке
     input_point = np.array([[point_x, point_y]])
     input_label = np.array([1])  # 1 = foreground
 
-    masks, scores, _ = sam_predictor.predict(
+    masks, scores, _ = sam_model.predict(
+        image=orig_np,
         point_coords=input_point,
         point_labels=input_label,
-        multimask_output=False  # Одна маска
+        multimask_output=False
     )
 
     # Берём лучшую маску
@@ -104,13 +103,13 @@ async def predict_full(
         result["freshness"] = fresh_name
         result["freshness_confidence"] = round(fresh_conf, 4)
 
-    # Возвращаем обрезанное изображение (по умолчанию 224×224)
+    # Возвращаем обрезанное изображение (по умолчанию cropped_size)
     if return_cropped:
-        cropped_final = crop_fruit_with_white_bg(orig_np, mask, out_size=224)
+        cropped_final = crop_fruit_with_white_bg(orig_np, mask, out_size=cropped_size)
         pil_img = Image.fromarray(cropped_final)
         buffered = io.BytesIO()
         pil_img.save(buffered, format="PNG")
         result["cropped_base64"] = base64.b64encode(buffered.getvalue()).decode('utf-8')
-        result["cropped_size"] = "224x224"
+        result["cropped_size"] = f"{cropped_size}x{cropped_size}"
 
     return result

models.py

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

utils.py

@@ -1,13 +1,11 @@
 import numpy as np
 import cv2
-import torch
 from PIL import Image
 import io
 import base64
+import torch
 from torchvision import transforms
-from mobile_sam import SamPredictor
 
-# Константы
 FRUIT_CLASSES = ['apple', 'banana', 'orange', 'strawberry', 'pear', 'lemon', 'cucumber', 'plum', 'raspberry', 'watermelon']
 FRESHNESS_CLASSES = ['freshapples', 'freshbanana', 'freshoranges', 'rottenapples', 'rottenbanana', 'rottenoranges']
 
@@ -19,7 +17,7 @@ def preprocess_for_classifier(img: np.ndarray) -> torch.Tensor:
     ])
     return transform(img)
 
-# Универсальный letterbox (без искажения)
+# Letterbox без искажения пропорций
 def letterbox_any_size(
     img: np.ndarray,
     target_size: int = 224,
@@ -45,11 +43,10 @@ def letterbox_any_size(
 # Обрезка по маске SAM + белый фон + letterbox
 def crop_fruit_with_white_bg(
     orig_img: np.ndarray,      # RGB
-    mask: np.ndarray,          # bool или uint8 от SAM
+    mask: np.ndarray,          # bool от SAM
     out_size: int = 224,
     bg_color: tuple = (255, 255, 255)
 ) -> np.ndarray:
-    # Маска → binary
     mask_bin = mask.astype(np.uint8)
 
     ys, xs = np.where(mask_bin == 1)
@@ -62,5 +59,4 @@ def crop_fruit_with_white_bg(
     cropped = orig_img[y1:y2+1, x1:x2+1].copy()
 
     final = letterbox_any_size(cropped, target_size=out_size, bg_color=bg_color)
-
     return final