one endpoint

app.py

@@ -1,13 +1,10 @@
-from fastapi import FastAPI, UploadFile, File, Form
+from fastapi import FastAPI, UploadFile, File, Query
 import torch
 from models import load_model1, load_model2, load_model3
 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,
@@ -20,119 +17,106 @@ import io
 
 app = FastAPI()
 
-# Загрузка моделей один раз при старте
-model1 = load_model1()  # сегментация → weights/model1.pth
-model2 = load_model2()  # классификатор → weights/model2.pth
-model3 = load_model3()
+# Глобальная загрузка моделей
+model1 = load_model1()   # segmentation
+model2 = load_model2()   # fruit type (10 классов)
+model3 = load_model3()   # freshness (6 классов)
 
 DEVICE = torch.device('cpu')
 
+# Классы, для которых делаем свежесть
+FRESHNESS_ELIGIBLE = {'apple', 'banana', 'orange'}
+
 
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
 
 
-@app.post("/predict1")
-async def predict1(file: UploadFile = File(...)):
+@app.post("/predict_full")
+async def predict_full(
+    file: UploadFile = File(...),
+    return_mask: bool = Query(default=False, description="Вернуть base64 маску сегментации?")
+):
     """
-    Сегментация фрукта → возвращает маску 256×256 в base64
+    Полный пайплайн:
+    1. Сегментация → маска
+    2. Если фрукт найден → классификация сорта
+    3. Если сорт в ['apple', 'banana', 'orange'] → классификация свежести
     """
+    # 1. Чтение изображения
     content = await file.read()
     image = Image.open(io.BytesIO(content)).convert('RGB')
-    image_np = np.array(image)
-
-    input_tensor = preprocess_image(image_np).unsqueeze(0).to(DEVICE)
+    orig_np = np.array(image)
 
+    # 2. Сегментация
+    input_tensor = preprocess_image(orig_np).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
         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)
-    }
-
-
-@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, margin_ratio=0.02,
-    target_size=100,
-    bg_color=(255, 255, 255))
-
-    # 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()
-
-    pred_idx = int(np.argmax(probs))
-    confidence = float(probs[pred_idx])
-
-    return {
-        "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(
-        original_np,
+    mask_256 = postprocess_mask(logits)  # (256, 256) float [0,1]
+
+    # Проверка: есть ли фрукт? (площадь > 5%)
+    fruit_area_ratio = np.mean(mask_256 > 0.5)
+    if fruit_area_ratio < 0.05:
+        return {
+            "status": "no_fruit_detected",
+            "fruit_area_ratio": round(fruit_area_ratio, 4),
+            "fruit": None,
+            "fruit_confidence": None,
+            "freshness": None,
+            "freshness_confidence": None,
+            "mask_256_base64": mask_to_base64(mask_256) if return_mask else None
+        }
+
+    # 3. Обрезание под модель сорта (100×100)
+    cropped_100 = apply_mask_and_crop_letterbox(
+        orig_np,
         mask_256,
-        margin_ratio=0.05,
-        target_size=224,
+        margin_ratio=0.02,
+        target_size=100,
         bg_color=(255, 255, 255)
     )
 
-    # Preprocess + inference
-    input_tensor = preprocess_for_freshness(cropped_224).unsqueeze(0).to(DEVICE)
-
+    # 4. Классификация сорта
+    input_tensor2 = preprocess_for_classifier(cropped_100).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])
+        logits2 = model2(input_tensor2)
+        probs2 = torch.softmax(logits2, dim=1).squeeze().cpu().numpy()
+
+    fruit_idx = int(np.argmax(probs2))
+    fruit_name = FRUIT_CLASSES[fruit_idx]
+    fruit_conf = float(probs2[fruit_idx])
+
+    result = {
+        "status": "success",
+        "fruit_area_ratio": round(fruit_area_ratio, 4),
+        "fruit": fruit_name,
+        "fruit_confidence": round(fruit_conf, 4),
+        "freshness": None,
+        "freshness_confidence": None,
+        "mask_256_base64": mask_to_base64(mask_256) if return_mask else None
+    }
 
-    return {
-        "predicted_class": FRESHNESS_CLASSES[pred_idx],
-        "confidence": round(confidence, 4),
-        "class_index": pred_idx
-    }
+    # 5. Если фрукт подходит — делаем свежесть
+    if fruit_name in FRESHNESS_ELIGIBLE:
+        cropped_224 = apply_mask_and_crop_letterbox(
+            orig_np,
+            mask_256,
+            margin_ratio=0.05,
+            target_size=224,
+            bg_color=(255, 255, 255)
+        )
+
+        input_tensor3 = preprocess_for_freshness(cropped_224).unsqueeze(0).to(DEVICE)
+        with torch.no_grad():
+            logits3 = model3(input_tensor3)
+            probs3 = torch.softmax(logits3, dim=1).squeeze().cpu().numpy()
+
+        fresh_idx = int(np.argmax(probs3))
+        fresh_name = FRESHNESS_CLASSES[fresh_idx]
+        fresh_conf = float(probs3[fresh_idx])
+
+        result["freshness"] = fresh_name
+        result["freshness_confidence"] = round(fresh_conf, 4)
+
+    return result