3rd endpoint

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, UploadFile, File, Form
 import torch
-from models import load_model1, load_model2
+from models import load_model1, load_model2, load_model3
 from utils import (
     # Для /predict1
     preprocess_image,
@@ -10,7 +10,10 @@ from utils import (
     decode_base64_mask,
     apply_mask_and_crop_letterbox,
     preprocess_for_classifier,
-    FRUIT_CLASSES
+    FRUIT_CLASSES,
+    apply_mask_and_crop_letterbox_224,
+    preprocess_for_freshness,
+    FRESHNESS_CLASSES
 )
 import numpy as np
 from PIL import Image
@@ -21,6 +24,7 @@ app = FastAPI()
 # Загрузка моделей один раз при старте
 model1 = load_model1()  # сегментация → weights/model1.pth
 model2 = load_model2()  # классификатор → weights/model2.pth
+model3 = load_model3()
 
 DEVICE = torch.device('cpu')
 
@@ -91,4 +95,42 @@ async def predict2(
         "predicted_fruit": FRUIT_CLASSES[pred_idx],
         "confidence": round(confidence, 4),
         "class_index": pred_idx
+    }
+
+
+@app.post("/predict3")
+async def predict3(
+    file: UploadFile = File(...),               # оригинальное изображение
+    mask_256_base64: str = Form(...)            # та же маска 256×256 от /predict1
+):
+    # Оригинал
+    content = await file.read()
+    original_pil = Image.open(io.BytesIO(content)).convert('RGB')
+    original_np = np.array(original_pil)
+
+    # Маска
+    mask_256 = decode_base64_mask(mask_256_base64)
+
+    # Вырезаем и готовим 224×224
+    cropped_224 = apply_mask_and_crop_letterbox_224(
+        original_np,
+        mask_256,
+        margin_ratio=0.05,          # подбери под свои тесты
+        bg_color=(255, 255, 255)
+    )
+
+    # Preprocess + inference
+    input_tensor = preprocess_for_freshness(cropped_224).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        logits = model3(input_tensor)
+        probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
+
+    pred_idx = int(np.argmax(probs))
+    confidence = float(probs[pred_idx])
+
+    return {
+        "predicted_class": FRESHNESS_CLASSES[pred_idx],
+        "confidence": round(confidence, 4),
+        "class_index": pred_idx
     }

models.py

@@ -5,7 +5,8 @@ import torch.nn as nn
 DEVICE = torch.device('cpu')
 
 model1 = None  # сегментация (как раньше)
-model2 = None  # классификатор
+model2 = None  # классификатор 1
+model3 = None  # классификатор 2
 
 def load_model1(weights_path='weights/seg.pth'):
     global model1
@@ -36,4 +37,18 @@ def load_model2(weights_path='weights/class1.pth'):
         model2.load_state_dict(state_dict)
         model2.to(DEVICE)
         model2.eval()
-    return model2
+    return model2
+
+
+def load_model3(weights_path='weights/class2.pth'):
+    global model3
+    if model3 is None:
+        model3 = models.mobilenet_v2(pretrained=False)
+        for param in model3.features.parameters():
+            param.requires_grad = False
+        model3.classifier[1] = nn.Linear(model3.classifier[1].in_features, 6)
+        state_dict = torch.load(weights_path, map_location=DEVICE)
+        model3.load_state_dict(state_dict)
+        model3.to(DEVICE)
+        model3.eval()
+    return model3

utils.py

@@ -130,4 +130,65 @@ def preprocess_for_classifier(img_100: np.ndarray) -> torch.Tensor:
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
-    return transform(img_100)
+    return transform(img_100)
+
+
+# ... весь предыдущий код остаётся ...
+
+# Новые константы для модели свежести
+FRESHNESS_CLASSES = [
+    'freshapples', 'freshbanana', 'freshoranges',
+    'rottenapples', 'rottenbanana', 'rottenoranges'
+]
+
+def preprocess_for_freshness(img_224: np.ndarray) -> torch.Tensor:
+    """ Трансформации, аналогичные test_transforms из обучения """
+    transform = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    return transform(img_224)
+
+
+def apply_mask_and_crop_letterbox_224(
+    orig_img: np.ndarray,
+    mask_256: np.ndarray,
+    margin_ratio: float = 0.05,           # можно подкрутить
+    bg_color: tuple = (255, 255, 255)     # белый фон — важно!
+) -> np.ndarray:
+    """
+    Аналог apply_mask_and_crop_letterbox, но для 224×224
+    """
+    # Letterbox до 224×224
+    letterbox_img, scale, paddings = letterbox_resize(orig_img, target_size=224)
+    top, bottom, left, right = paddings
+
+    # Применяем маску (маска 256→ресайзим до 224)
+    mask_resized = cv2.resize(mask_256, (224, 224), interpolation=cv2.INTER_NEAREST)
+
+    masked = letterbox_img.copy()
+    masked[mask_resized < 0.5] = bg_color   # белый фон
+
+    # Контуры
+    mask_bin = (mask_resized > 0.5).astype(np.uint8) * 255
+    contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    if not contours:
+        return np.full((224, 224, 3), bg_color, dtype=np.uint8)
+
+    cnt = max(contours, key=cv2.contourArea)
+    x, y, bw, bh = cv2.boundingRect(cnt)
+
+    margin = int(max(bw, bh) * margin_ratio)
+    x1 = max(0, x - margin)
+    y1 = max(0, y - margin)
+    x2 = min(224, x + bw + margin)
+    y2 = min(224, y + bh + margin)
+
+    cropped = masked[y1:y2, x1:x2]
+
+    # Финальный resize до 224×224 (если обрезали меньше)
+    final = cv2.resize(cropped, (224, 224), interpolation=cv2.INTER_AREA)
+
+    return final