mask fix

app.py

@@ -25,7 +25,7 @@ model3 = load_model3()   # freshness (6 классов)
 DEVICE = torch.device('cpu')
 
 # Классы, для которых делаем свежесть
-FRESHNESS_ELIGIBLE = {'apple', 'banana', 'orange'}
+FRESHNESS_ELIGIBLE = {'apple', 'banana', 'orange', 'lemon'}
 
 
 @app.get("/")
@@ -72,7 +72,7 @@ async def predict_full(
     cropped_100 = apply_mask_and_crop_letterbox(
         orig_np,
         mask_256,
-        margin_ratio=0.02,
+        margin_ratio=0.01,
         target_size=100,
         bg_color=(255, 255, 255)
     )
@@ -102,7 +102,7 @@ async def predict_full(
         cropped_224 = apply_mask_and_crop_letterbox(
             orig_np,
             mask_256,
-            margin_ratio=0.05,
+            margin_ratio=0.03,
             target_size=224,
             bg_color=(255, 255, 255)
         )

utils.py

@@ -19,19 +19,18 @@ def preprocess_image(image_np: np.ndarray) -> torch.Tensor:
     augmented = preprocess_transform(image=image_np)
     return augmented['image']
 
-def postprocess_mask(logits: torch.Tensor, threshold: float = 0.5) -> np.ndarray:
+def postprocess_mask(logits: torch.Tensor, threshold: float = 0.7) -> np.ndarray:
     pred = torch.sigmoid(logits).squeeze().cpu().numpy()
     binary_mask = (pred > threshold).astype(np.float32)
+    kernel = np.ones((3, 3), np.uint8)
+    mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)  # заполнить дырки
+    binary_mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)  # убрать шум
     return binary_mask  # shape (256, 256)
 
 # ────────────────────────────────────────────────
 # Для /predict1 — возвращаем маску 256×256
 # ────────────────────────────────────────────────
 
-def resize_mask(mask: np.ndarray, size: int = 256) -> np.ndarray:
-    return cv2.resize(mask, (size, size), interpolation=cv2.INTER_NEAREST)
-
-
 def mask_to_base64(mask: np.ndarray) -> str:
     pil_mask = Image.fromarray((mask * 255).astype(np.uint8)).convert('L')
     buffered = io.BytesIO()
@@ -47,13 +46,6 @@ def mask_to_base64(mask: np.ndarray) -> str:
 FRUIT_CLASSES = ['apple', 'banana', 'orange', 'strawberry', 'pear', 'lemon', 'cucumber', 'plum', 'raspberry', 'watermelon']
 
 
-def decode_base64_mask(base64_str: str) -> np.ndarray:
-    img_data = base64.b64decode(base64_str)
-    pil_img = Image.open(io.BytesIO(img_data)).convert('L')
-    mask = np.array(pil_img) / 255.0
-    return mask.astype(np.float32)  # shape ≈ (256, 256)
-
-
 def letterbox_resize(img: np.ndarray, target_size: int = 256) -> tuple[np.ndarray, float, tuple[int, int]]:
     """
     Resize с сохранением пропорций + padding чёрным