Sumeshbuilds/specimen-backend

1import io
2import os
3import base64
4import collections
5import numpy as np
6from PIL import Image
7from ultralytics import YOLO
8from huggingface_hub import hf_hub_download
9from fastapi import File, UploadFile
10from fastapi.middleware.cors import CORSMiddleware
11from fastapi.responses import JSONResponse
12import gradio as gr
13import spaces  # Import HF ZeroGPU scheduling
14
15# =====================================================================
16# STEP 1: DYNAMICALLY DOWNLOAD ALL 3 ENGINES FROM KORTIV AI
17# =====================================================================
18try:
19    print("Downloading Model 1: Plant Segmenter...")
20    seg_path = hf_hub_download(repo_id="kortivai/SpecimenAI-PlantDoc", filename="plant_segmenter.pt")
21    segmenter = YOLO(seg_path)
22
23    print("Downloading Model 2: Label Detector...")
24    lbl_path = hf_hub_download(repo_id="kortivai/SpecimenAI-PlantDoc", filename="label_detector.pt")
25    label_detector = YOLO(lbl_path)
26
27    print("Downloading Model 3: Organ Counter...")
28    org_path = hf_hub_download(repo_id="kortivai/SpecimenAI-PlantDoc", filename="organ_counter.pt")
29    organ_counter = YOLO(org_path)
30
31    print("✔ All 3 YOLO11/YOLO26-s engines loaded successfully into memory!")
32except Exception as e:
33    print(f"Error loading models: {e}")
34    segmenter, label_detector, organ_counter = None, None, None
35
36# =====================================================================
37# STEP 2: THE CHAINED PIPELINE ENGINE (RUNS ON A100 GPU)
38# =====================================================================
39@spaces.GPU
40def run_scientific_pipeline(img_array):
41    if img_array is None:
42        return None, "", 0.0, {}
43
44    # --- PIPELINE STAGE 1: DETECT & CROP THE PAPER LABEL ---
45    cropped_label_b64 = ""
46    label_results = label_detector.predict(source=img_array, conf=0.25, device="cuda", save=False)
47    label_boxes = label_results[0].boxes
48    
49    if len(label_boxes) > 0:
50        # Take the label box with the highest confidence
51        best_idx = label_boxes.conf.cpu().numpy().argmax()
52        xyxy = label_boxes.xyxy[best_idx].cpu().numpy().astype(int)
53        x1, y1, x2, y2 = xyxy
54        
55        # Crop the label out of the raw image array
56        cropped_label = img_array[y1:y2, x1:x2]
57        
58        # Convert cropped array to PIL -> Base64
59        cropped_pil = Image.fromarray(cropped_label)
60        buffered = io.BytesIO()
61        cropped_pil.save(buffered, format="JPEG")
62        cropped_label_b64 = "data:image/jpeg;base64," + base64.b64encode(buffered.getvalue()).decode("utf-8")
63
64    # --- PIPELINE STAGE 2: CALCULATE PLANT COVERAGE (%) ---
65    plant_coverage_pct = 0.0
66    seg_results = segmenter.predict(source=img_array, conf=0.25, device="cuda", save=False)
67    plant_boxes = seg_results[0].boxes
68    
69    if len(plant_boxes) > 0:
70        total_pixels = img_array.shape[0] * img_array.shape[1]
71        plant_pixels = 0
72        # Calculate sum of detected specimen bounding box areas
73        for box in plant_boxes.xyxy.cpu().numpy():
74            x1, y1, x2, y2 = box
75            plant_pixels += (x2 - x1) * (y2 - y1)
76        plant_coverage_pct = min(100.0, (plant_pixels / total_pixels) * 100.0)
77
78    # --- PIPELINE STAGE 3: DETECT & COUNT BIOLOGICAL ORGANS ---
79    organ_results = organ_counter.predict(source=img_array, conf=0.25, device="cuda", save=False)
80    
81    # Draw organ detections (Leaf, Flower, Fruit)
82    annotated_img = organ_results[0].plot()
83    annotated_img_rgb = annotated_img[:, :, ::-1] # BGR to RGB
84    pil_annotated = Image.fromarray(annotated_img_rgb)
85    
86    # Count organs
87    organ_boxes = organ_results[0].boxes
88    counts = {}
89    if len(organ_boxes) > 0:
90        class_ids = organ_boxes.cls.cpu().numpy().astype(int)
91        class_names = [organ_counter.names[cid] for cid in class_ids]
92        counts = dict(collections.Counter(class_names))
93
94    return pil_annotated, cropped_label_b64, round(plant_coverage_pct, 2), counts
95
96# =====================================================================
97# STEP 3: DEFINE GRADIO INTERFACE
98# =====================================================================
99demo = gr.Interface(
100    fn=run_scientific_pipeline,
101    inputs=gr.Image(type="numpy", label="Upload Specimen Sheet"),
102    outputs=[
103        gr.Image(type="pil", label="Annotated Specimen"),
104        gr.Textbox(label="Cropped Label (Base64)"),
105        gr.Number(label="Plant Coverage (%)"),
106        gr.JSON(label="Detected Metrics")
107    ],
108    title="🌱 Kortiv AI - SpecimenAI API Portal",
109    api_name="predict_leaf"
110)
111
112# =====================================================================
113# STEP 4: FASTAPI CORS OVERRIDES (FOR DECOUPLED VERCEL APP)
114# =====================================================================
115app = demo.app
116
117app.add_middleware(
118    CORSMiddleware,
119    allow_origins=["*"],
120    allow_credentials=True,
121    allow_methods=["*"],
122    allow_headers=["*"],
123)
124
125# Injected route for direct REST API queries (if required)
126@app.post("/api/diagnose")
127async def predict_api(file: UploadFile = File(...)):
128    if not segmenter or not label_detector or not organ_counter:
129        return JSONResponse(status_code=500, content={"error": "One or more pipeline models failed to load"})
130    
131    contents = await file.read()
132    image = Image.open(io.BytesIO(contents)).convert("RGB")
133    img_array = np.array(image)
134
135    # Execute the GPU-accelerated pipeline
136    pil_annotated, cropped_label_b64, plant_coverage_pct, counts = run_scientific_pipeline(img_array)
137    
138    # Convert annotated image to Base64
139    buffered = io.BytesIO()
140    pil_annotated.save(buffered, format="JPEG")
141    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
142
143    return {
144        "image": f"data:image/jpeg;base64,{img_str}",
145        "cropped_label": cropped_label_b64,
146        "plant_coverage": plant_coverage_pct,
147        "metrics": counts
148    }
149
150# Launch the server globally
151demo.launch()