mobilesam v1.5 (test new mask)

app.py

@@ -6,7 +6,7 @@ import io
 import base64
 from models import load_sam, load_model2, load_model3
 from utils import (
-    crop_fruit_with_white_bg,
+    crop_fruit_contour_letterbox,
     preprocess_for_classifier,
     FRUIT_CLASSES,
     FRESHNESS_CLASSES
@@ -71,7 +71,7 @@ async def predict_full(
         }
 
     # Обрезка под 100×100 для сорта
-    cropped_100 = crop_fruit_with_white_bg(orig_np, mask, out_size=100)
+    cropped_100 = crop_fruit_contour_letterbox(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)
@@ -93,7 +93,7 @@ async def predict_full(
 
     # Свежесть, если подходит
     if fruit_name in FRESHNESS_ELIGIBLE:
-        cropped_224 = crop_fruit_with_white_bg(orig_np, mask, out_size=224)
+        cropped_224 = crop_fruit_contour_letterbox(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)
@@ -108,7 +108,7 @@ async def predict_full(
 
     # Возвращаем обрезанное изображение
     if return_cropped:
-        cropped_final = crop_fruit_with_white_bg(orig_np, mask, out_size=cropped_size)
+        cropped_final = crop_fruit_contour_letterbox(orig_np, mask, out_size=cropped_size)
         pil_img = Image.fromarray(cropped_final)
         buffered = io.BytesIO()
         pil_img.save(buffered, format="PNG")

utils.py

@@ -39,13 +39,13 @@ def letterbox_any_size(
                                 cv2.BORDER_CONSTANT, value=bg_color)
     return padded
 
-def crop_fruit_with_white_bg(
+def crop_fruit_contour_letterbox(
     orig_img: np.ndarray,
     mask: np.ndarray,
     out_size: int = 224,
     bg_color: tuple = (255, 255, 255)
 ) -> np.ndarray:
-    mask_bin = mask.astype(np.uint8)
+    mask_bin = (mask > 0.5).astype(np.uint8)
 
     ys, xs = np.where(mask_bin == 1)
     if len(xs) == 0:
@@ -54,7 +54,17 @@ def crop_fruit_with_white_bg(
     y1, y2 = ys.min(), ys.max()
     x1, x2 = xs.min(), xs.max()
 
-    cropped = orig_img[y1:y2+1, x1:x2+1].copy()
+    cropped_rgb = orig_img[y1:y2+1, x1:x2+1].copy()
+    cropped_mask = mask_bin[y1:y2+1, x1:x2+1]
+
+    # Сначала letterbox
+    letterboxed = letterbox_any_size(cropped_rgb, target_size=out_size, bg_color=bg_color)
+
+    # Масштабируем маску под letterbox (это сложнее, но можно приблизить)
+    # Для простоты — применяем маску на cropped перед letterbox
+    white_bg = np.full_like(cropped_rgb, bg_color)
+    masked_cropped = np.where(cropped_mask[..., None] == 1, cropped_rgb, white_bg)
+
+    final = letterbox_any_size(masked_cropped, target_size=out_size, bg_color=bg_color)
 
-    final = letterbox_any_size(cropped, target_size=out_size, bg_color=bg_color)
     return final