mobilesam v1.0

app.py

@@ -1,30 +1,26 @@
 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
 from utils import (
-    preprocess_image,
-    predict_mask_tta,
-    postprocess_mask,
-    mask_to_base64,
-    apply_white_background_and_crop,
+    crop_fruit_with_white_bg,
     preprocess_for_classifier,
     FRUIT_CLASSES,
     FRESHNESS_CLASSES
 )
-import numpy as np
-from PIL import Image
-import io
 
 app = FastAPI()
 
-# Глобальная загрузка моделей
-model1 = load_model1()   # segmentation (448)
-model2 = load_model2()   # fruit type
-model3 = load_model3()   # freshness
+# Загрузка моделей
+sam_predictor = load_model1()  # MobileSAM
+model2 = load_model2()
+model3 = load_model3()
 
 DEVICE = torch.device('cpu')
 
-# Классы, для которых делаем свежесть
 FRESHNESS_ELIGIBLE = {'apple', 'banana', 'orange', 'lemon'}
 
 @app.get("/")
@@ -34,19 +30,33 @@ def greet_json():
 @app.post("/predict_full")
 async def predict_full(
     file: UploadFile = File(...),
-    return_mask: bool = Query(default=False, description="Вернуть base64 маски?")
+    point_x: int = Query(..., description="X-координата точки на фрукте (в пикселях оригинального изображения)"),
+    point_y: int = Query(..., description="Y-координата точки на фрукте"),
+    return_cropped: bool = Query(default=True, description="Вернуть обрезанное изображение в base64?")
 ):
     content = await file.read()
     image = Image.open(io.BytesIO(content)).convert('RGB')
     orig_np = np.array(image)
 
-    # Сегментация
-    input_tensor = preprocess_image(orig_np).unsqueeze(0).to(DEVICE)
-    with torch.no_grad():
-        prob = predict_mask_tta(model1, input_tensor)
-    mask = postprocess_mask(prob.squeeze().cpu().numpy())
+    # Установка изображения в SAM
+    sam_predictor.set_image(orig_np)
 
-    fruit_area_ratio = np.mean(mask > 0.5)
+    # Промпт: точка на фрукте
+    input_point = np.array([[point_x, point_y]])
+    input_label = np.array([1])  # 1 = foreground
+
+    masks, scores, _ = sam_predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=False  # Одна маска
+    )
+
+    # Берём лучшую маску
+    best_mask_idx = np.argmax(scores)
+    mask = masks[best_mask_idx]  # bool
+
+    # Проверка: есть ли фрукт?
+    fruit_area_ratio = np.mean(mask)
     if fruit_area_ratio < 0.01:
         return {
             "status": "no_fruit_detected",
@@ -55,12 +65,11 @@ async def predict_full(
             "fruit_confidence": None,
             "freshness": None,
             "freshness_confidence": None,
-            "mask_base64": mask_to_base64(mask) if return_mask else None
+            "cropped_base64": None
         }
 
-    # Для сорта (100×100)
-    cropped_100 = apply_white_background_and_crop(orig_np, mask, out_size=100)
-
+    # Обрезка под 100×100 для сорта
+    cropped_100 = crop_fruit_with_white_bg(orig_np, mask, out_size=100)
     input_tensor2 = preprocess_for_classifier(cropped_100).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
         logits2 = model2(input_tensor2)
@@ -77,13 +86,12 @@ async def predict_full(
         "fruit_confidence": round(fruit_conf, 4),
         "freshness": None,
         "freshness_confidence": None,
-        "mask_base64": mask_to_base64(mask) if return_mask else None
+        "cropped_base64": None
     }
 
-    # Если фрукт подходит — свежесть (224×224)
+    # Свежесть, если подходит
     if fruit_name in FRESHNESS_ELIGIBLE:
-        cropped_224 = apply_white_background_and_crop(orig_np, mask, out_size=224)
-
+        cropped_224 = crop_fruit_with_white_bg(orig_np, mask, out_size=224)
         input_tensor3 = preprocess_for_classifier(cropped_224).unsqueeze(0).to(DEVICE)
         with torch.no_grad():
             logits3 = model3(input_tensor3)
@@ -96,4 +104,13 @@ async def predict_full(
         result["freshness"] = fresh_name
         result["freshness_confidence"] = round(fresh_conf, 4)
 
+    # Возвращаем обрезанное изображение (по умолчанию 224×224)
+    if return_cropped:
+        cropped_final = crop_fruit_with_white_bg(orig_np, mask, out_size=224)
+        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"
+
     return result

models.py

@@ -2,25 +2,20 @@ 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
 
 DEVICE = torch.device('cpu')
 
-model1 = None  # сегментация
+model1 = None  # теперь это MobileSAM
 model2 = None  # сорт фрукта
 model3 = None  # свежесть
 
-def load_model1(weights_path='weights/seg0.pth'):
+def load_model1(weights_path='weights/mobile_sam.pt'):
     global model1
     if model1 is None:
-        model1 = smp.Unet(
-            encoder_name="mobilenet_v2",
-            encoder_weights=None,
-            in_channels=3,
-            classes=1,
-            activation=None
-        ).to(DEVICE)
-        state_dict = torch.load(weights_path, map_location=DEVICE)
-        model1.load_state_dict(state_dict)
+        model_type = "vit_t"
+        model1 = sam_model_registry[model_type](checkpoint=weights_path)
+        model1.to(DEVICE)
         model1.eval()
     return model1
 

requirements.txt

@@ -7,4 +7,5 @@ albumentations
 pillow
 numpy
 opencv-python-headless
-python-multipart
+python-multipart
+git+https://github.com/ChaoningZhang/MobileSAM.git

utils.py

@@ -1,100 +1,16 @@
 import numpy as np
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-import torch
 import cv2
+import torch
 from PIL import Image
 import io
 import base64
 from torchvision import transforms
+from mobile_sam import SamPredictor
 
-# ────────────────────────────────────────────────
-# Новый размер входа модели — 448×448
-# ────────────────────────────────────────────────
-IMG_SIZE = 448
-
-preprocess_transform = A.Compose([
-    A.Resize(IMG_SIZE, IMG_SIZE),
-    A.Normalize(),  # mean/std ImageNet — то же, что и в обучении
-    ToTensorV2()
-])
-
-def preprocess_image(image_np: np.ndarray) -> torch.Tensor:
-    augmented = preprocess_transform(image=image_np)
-    return augmented['image']
-
-# ────────────────────────────────────────────────
-# TTA-предсказание (как в твоём примере)
-# ────────────────────────────────────────────────
-@torch.no_grad()
-def predict_mask_tta(model, image_tensor):
-    preds = []
-    # Оригинал
-    preds.append(torch.sigmoid(model(image_tensor)))
-    # Flip horizontal
-    preds.append(
-        torch.flip(
-            torch.sigmoid(model(torch.flip(image_tensor, dims=[3]))),
-            dims=[3]
-        )
-    )
-    # Flip vertical
-    preds.append(
-        torch.flip(
-            torch.sigmoid(model(torch.flip(image_tensor, dims=[2]))),
-            dims=[2]
-        )
-    )
-    return torch.mean(torch.stack(preds), dim=0)
-
-# ────────────────────────────────────────────────
-# Post-processing маски (как в твоём примере + морфология)
-# ────────────────────────────────────────────────
-def postprocess_mask(prob: np.ndarray, threshold: float = 0.65, min_area_ratio: float = 0.01) -> np.ndarray:
-    binary = (prob > threshold).astype(np.uint8)
-
-    # Connected components — оставляем только главный объект
-    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
-
-    if num_labels <= 1:
-        return binary.astype(np.float32)
-
-    largest_label = np.argmax(stats[1:, cv2.CC_STAT_AREA]) + 1
-    area = stats[largest_label, cv2.CC_STAT_AREA]
-
-    if area < binary.shape[0] * binary.shape[1] * min_area_ratio:
-        return np.zeros_like(binary, dtype=np.float32)
-
-    clean_mask = (labels == largest_label).astype(np.float32)
-
-    # Морфология (заполнить дырки, убрать шум)
-    kernel = np.ones((3, 3), np.uint8)
-    clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_CLOSE, kernel)
-    clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_OPEN, kernel)
-
-    return clean_mask
-
-# ────────────────────────────────────────────────
-# Base64 маски (для возврата клиенту)
-# ────────────────────────────────────────────────
-def mask_to_base64(mask: np.ndarray) -> str:
-    pil_mask = Image.fromarray((mask * 255).astype(np.uint8)).convert('L')
-    buffered = io.BytesIO()
-    pil_mask.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-# ────────────────────────────────────────────────
-# Константы классов
-# ────────────────────────────────────────────────
+# Константы
 FRUIT_CLASSES = ['apple', 'banana', 'orange', 'strawberry', 'pear', 'lemon', 'cucumber', 'plum', 'raspberry', 'watermelon']
-FRESHNESS_CLASSES = [
-    'freshapples', 'freshbanana', 'freshoranges',
-    'rottenapples', 'rottenbanana', 'rottenoranges'
-]
+FRESHNESS_CLASSES = ['freshapples', 'freshbanana', 'freshoranges', 'rottenapples', 'rottenbanana', 'rottenoranges']
 
-# ────────────────────────────────────────────────
-# Preprocess для классификаторов (100 и 224)
-# ────────────────────────────────────────────────
 def preprocess_for_classifier(img: np.ndarray) -> torch.Tensor:
     transform = transforms.Compose([
         transforms.ToPILImage(),
@@ -103,9 +19,7 @@ def preprocess_for_classifier(img: np.ndarray) -> torch.Tensor:
     ])
     return transform(img)
 
-# ────────────────────────────────────────────────
-# Универсальный letterbox (для любого target_size)
-# ────────────────────────────────────────────────
+# Универсальный letterbox (без искажения)
 def letterbox_any_size(
     img: np.ndarray,
     target_size: int = 224,
@@ -124,22 +38,19 @@ def letterbox_any_size(
     left = pad_w // 2
     right = pad_w - left
 
-    padded = cv2.copyMakeBorder(
-        resized, top, bottom, left, right,
-        cv2.BORDER_CONSTANT, value=bg_color
-    )
+    padded = cv2.copyMakeBorder(resized, top, bottom, left, right,
+                                cv2.BORDER_CONSTANT, value=bg_color)
     return padded
 
-# ────────────────────────────────────────────────
-# Apply white background + crop по маске (448×448)
-# ────────────────────────────────────────────────
-def apply_white_background_and_crop(
-    orig_img: np.ndarray,          # RGB
-    mask: np.ndarray,              # float [0,1] 448×448
+# Обрезка по маске SAM + белый фон + letterbox
+def crop_fruit_with_white_bg(
+    orig_img: np.ndarray,      # RGB
+    mask: np.ndarray,          # bool или uint8 от SAM
     out_size: int = 224,
     bg_color: tuple = (255, 255, 255)
 ) -> np.ndarray:
-    mask_bin = (mask > 0.5).astype(np.uint8)
+    # Маска → binary
+    mask_bin = mask.astype(np.uint8)
 
     ys, xs = np.where(mask_bin == 1)
     if len(xs) == 0:

weights/{seg0.pth → mobile_sam.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4e2a778652280420b80ebc949fac8e2d1a95737d28884e6fa99df2509c7410db
-size 26806811
+oid sha256:6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f
+size 40728226