add class1

app.py

@@ -1,15 +1,28 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, Form
 import torch
-from models import load_model1
-from utils import preprocess_image, postprocess_mask, resize_mask, mask_to_base64
+from models import load_model1, load_model2
+from utils import (
+    # Для /predict1
+    preprocess_image,
+    postprocess_mask,
+    mask_to_base64,
+    # Для /predict2
+    decode_base64_mask,
+    apply_mask_and_crop_letterbox,
+    preprocess_for_classifier,
+    FRUIT_CLASSES
+)
 import numpy as np
 from PIL import Image
 import io
 
 app = FastAPI()
 
-# Загрузка модели при старте (глобально, один раз)
-model1 = load_model1()
+# Загрузка моделей один раз при старте
+model1 = load_model1()  # сегментация → weights/model1.pth
+model2 = load_model2()  # классификатор → weights/model2.pth
+
+DEVICE = torch.device('cpu')
 
 
 @app.get("/")
@@ -19,20 +32,63 @@ def greet_json():
 
 @app.post("/predict1")
 async def predict1(file: UploadFile = File(...)):
+    """
+    Сегментация фрукта → возвращает маску 256×256 в base64
+    """
     content = await file.read()
     image = Image.open(io.BytesIO(content)).convert('RGB')
     image_np = np.array(image)
 
-    input_tensor = preprocess_image(image_np)
+    input_tensor = preprocess_image(image_np).unsqueeze(0).to(DEVICE)
+
     with torch.no_grad():
-        logits = model1(input_tensor.unsqueeze(0))  # batch dim
+        logits = model1(input_tensor)
+
+    pred_mask = postprocess_mask(logits)  # shape (256, 256), float [0,1]
+
+    # Возвращаем только одну маску — 256×256
+    return {
+        "mask_256_base64": mask_to_base64(pred_mask)
+    }
+
 
-    pred_mask = postprocess_mask(logits)  # (256, 256) binary
+@app.post("/predict2")
+async def predict2(
+    file: UploadFile = File(...),               # оригинальное изображение (любого размера)
+    mask_256_base64: str = Form(...)            # маска 256×256 от /predict1
+):
+    """
+    Классификация фрукта:
+    - ресайз оригинала → letterbox 256×256
+    - применение маски
+    - crop по bounding box + margin
+    - ресайз результата до 100×100
+    - inference MobileNetV2
+    """
+    # 1. Оригинальное изображение
+    content = await file.read()
+    original_pil = Image.open(io.BytesIO(content)).convert('RGB')
+    original_np = np.array(original_pil)
+
+    # 2. Декодируем маску 256×256
+    mask_256 = decode_base64_mask(mask_256_base64)
+
+    # 3. Letterbox + маска + crop + resize до 100×100
+    cropped_100 = apply_mask_and_crop_letterbox(original_np, mask_256)
+
+    # 4. Препроцессинг для классификатора
+    input_tensor = preprocess_for_classifier(cropped_100).unsqueeze(0).to(DEVICE)
+
+    # 5. Инференс
+    with torch.no_grad():
+        logits = model2(input_tensor)
+        probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
 
-    mask_100 = resize_mask(pred_mask, 100)
-    mask_224 = resize_mask(pred_mask, 224)
+    pred_idx = int(np.argmax(probs))
+    confidence = float(probs[pred_idx])
 
     return {
-        "mask_100_base64": mask_to_base64(mask_100),
-        "mask_224_base64": mask_to_base64(mask_224)
+        "predicted_fruit": FRUIT_CLASSES[pred_idx],
+        "confidence": round(confidence, 4),
+        "class_index": pred_idx
     }

models.py

@@ -1,13 +1,16 @@
 import torch
-import segmentation_models_pytorch as smp
+import torchvision.models as models
+import torch.nn as nn
 
 DEVICE = torch.device('cpu')
 
-model1 = None
+model1 = None  # сегментация (как раньше)
+model2 = None  # классификатор
 
 def load_model1(weights_path='weights/seg.pth'):
     global model1
     if model1 is None:
+        import segmentation_models_pytorch as smp
         model1 = smp.Unet(
             encoder_name="mobilenet_v2",
             encoder_weights=None,
@@ -17,4 +20,20 @@ def load_model1(weights_path='weights/seg.pth'):
         state_dict = torch.load(weights_path, map_location=DEVICE)
         model1.load_state_dict(state_dict)
         model1.eval()
-    return model1
+    return model1
+
+
+def load_model2(weights_path='weights/class1.pth'):
+    global model2
+    if model2 is None:
+        model2 = models.mobilenet_v2(pretrained=False)
+        # Замораживаем features (как в обучении)
+        for param in model2.features.parameters():
+            param.requires_grad = False
+        # Заменяем classifier
+        model2.classifier[1] = nn.Linear(model2.classifier[1].in_features, 10)
+        state_dict = torch.load(weights_path, map_location=DEVICE)
+        model2.load_state_dict(state_dict)
+        model2.to(DEVICE)
+        model2.eval()
+    return model2

requirements.txt

@@ -1,8 +1,10 @@
 fastapi
 uvicorn[standard]
 torch
+torchvision
 segmentation_models_pytorch
 albumentations
 pillow
 numpy
-opencv-python-headless
+opencv-python-headless
+python-multipart

utils.py

@@ -6,6 +6,7 @@ import cv2
 from PIL import Image
 import io
 import base64
+from torchvision import transforms
 
 # Препроцессинг: аналог валидации
 preprocess_transform = A.Compose([
@@ -23,14 +24,110 @@ def postprocess_mask(logits: torch.Tensor, threshold: float = 0.5) -> np.ndarray
     binary_mask = (pred > threshold).astype(np.float32)
     return binary_mask  # shape (256, 256)
 
-def resize_mask(mask: np.ndarray, size: int) -> np.ndarray:
-    # Resize с nearest neighbor для бинарных масок
-    resized = cv2.resize(mask, (size, size), interpolation=cv2.INTER_NEAREST)
-    return resized.astype(np.float32)  # 0/1 float
+# ────────────────────────────────────────────────
+# Для /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 grayscale (0/255), save as PNG, base64
     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')
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+
+# ────────────────────────────────────────────────
+# Для /predict2
+# ────────────────────────────────────────────────
+
+# Новые для классификации
+FRUIT_CLASSES = ['apple', 'banana', 'orange', 'grape', 'strawberry',
+                 'tomato', 'pear', 'peach', 'cherry', 'lemon']
+
+
+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 чёрным
+    Возвращает: новое изображение, scale_factor, (pad_top, pad_bottom, pad_left, pad_right)
+    """
+    h, w = img.shape[:2]
+    scale = min(target_size / h, target_size / w)
+    new_h, new_w = int(h * scale), int(w * scale)
+
+    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+
+    pad_h = target_size - new_h
+    pad_w = target_size - new_w
+    top = pad_h // 2
+    bottom = pad_h - top
+    left = pad_w // 2
+    right = pad_w - left
+
+    padded = cv2.copyMakeBorder(
+        resized, top, bottom, left, right,
+        cv2.BORDER_CONSTANT, value=(0, 0, 0)
+    )
+
+    return padded, scale, (top, bottom, left, right)
+
+
+def apply_mask_and_crop_letterbox(
+    orig_img: np.ndarray,          # оригинал любой размер
+    mask_256: np.ndarray           # маска 256×256 [0..1]
+) -> np.ndarray:
+    """
+    1. Делаем letterbox-версию оригинала 256×256
+    2. Применяем маску
+    3. Находим bbox
+    4. Вырезаем + margin
+    5. Ресайзим до 100×100
+    """
+    letterbox_img, scale, paddings = letterbox_resize(orig_img, 256)
+    top, bottom, left, right = paddings
+
+    # Маска уже 256×256 — применяем напрямую
+    masked = letterbox_img.copy()
+    masked[mask_256 < 0.5] = 0
+
+    # Находим контуры / bbox
+    mask_bin = (mask_256 > 0.5).astype(np.uint8) * 255
+    contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    if not contours:
+        return np.zeros((100, 100, 3), dtype=np.uint8)
+
+    cnt = max(contours, key=cv2.contourArea)
+    x, y, bw, bh = cv2.boundingRect(cnt)
+
+    # margin ~10%
+    margin = int(max(bw, bh) * 0.12)
+    x1 = max(0, x - margin)
+    y1 = max(0, y - margin)
+    x2 = min(256, x + bw + margin)
+    y2 = min(256, y + bh + margin)
+
+    cropped = masked[y1:y2, x1:x2]
+
+    # Финальный ресайз до 100×100 для классификатора
+    final = cv2.resize(cropped, (100, 100), interpolation=cv2.INTER_AREA)
+
+    return final
+
+
+def preprocess_for_classifier(img_100: np.ndarray) -> torch.Tensor:
+    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_100)