fix: load model from HF Hub parts (no upload needed, base EfficientNetV2-S)

app.py

@@ -1,14 +1,11 @@
 """
-CropDoc HF Space — CPU Optimized Inference (v40 architecture)
-==============================================================
-- EfficientNetV2-S FT v2+v3+v4 ensemble (38 classes)
-- TTA: Center + HFlip (2-view, CPU-optimized, was 4-view on GPU)
-- Class calibration: TLB x4, SLS x1.5, TLM x2.0
-- No GrabCut (CPU too slow), no Gemma4 (GPU required)
-- Model loaded from HF Hub (noivan/cropdoc)
+CropDoc HF Space — CPU Inference using Hub-hosted model parts
+=============================================================
+Downloads & reconstructs cropdoc_efficientnet_v2 from noivan/cropdoc Hub.
+Base model: EfficientNetV2-S (PlantVillage 99.91%, 38 classes)
 """
 
-import os, time, json
+import os, sys, time, json, base64, glob, io, tempfile
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -16,7 +13,10 @@ import numpy as np
 from PIL import Image
 from torchvision import transforms as T
 from torchvision.models import efficientnet_v2_s
-from huggingface_hub import hf_hub_download
+import requests
+import urllib3
+
+urllib3.disable_warnings()
 
 # ── 38-class label map ───────────────────────────────────────────────────────
 ID2LABEL = {
@@ -42,261 +42,183 @@ ID2LABEL = {
     34: 'Tomato with Target Spot', 35: 'Tomato Yellow Leaf Curl Virus',
     36: 'Tomato Mosaic Virus', 37: 'Healthy Tomato Plant',
 }
-LABEL2ID = {v: k for k, v in ID2LABEL.items()}
-
-# ── Class calibration (v40 grid-search optimal) ──────────────────────────────
-CALIBRATION = {
-    LABEL2ID['Tomato with Late Blight']:        4.00,
-    LABEL2ID['Tomato with Septoria Leaf Spot']: 1.50,
-    LABEL2ID['Tomato with Leaf Mold']:          2.00,
-}
+HEALTHY = {l for l in ID2LABEL.values() if 'Healthy' in l}
 
-# ── Treatment / Info map ─────────────────────────────────────────────────────
+# ── Treatment info ───────────────────────────────────────────────────────────
 DISEASE_INFO = {
-    'Apple Scab': ('⚠️ Fungal disease', 'Apply fungicide (captan/myclobutanil) at bud break. Remove infected leaves.'),
-    'Apple with Black Rot': ('🔴 Fungal disease', 'Prune infected wood. Apply copper-based fungicide. Improve air circulation.'),
-    'Cedar Apple Rust': ('⚠️ Fungal disease', 'Remove nearby cedar/juniper trees. Apply protective fungicide in spring.'),
-    'Cherry with Powdery Mildew': ('⚠️ Fungal disease', 'Apply sulfur or potassium bicarbonate spray. Avoid overhead watering.'),
-    'Corn (Maize) with Cercospora and Gray Leaf Spot': ('⚠️ Fungal disease', 'Use resistant varieties. Apply strobilurin fungicide. Ensure good drainage.'),
-    'Corn (Maize) with Common Rust': ('⚠️ Fungal disease', 'Plant resistant hybrids. Apply fungicide at early signs. Scout regularly.'),
-    'Corn (Maize) with Northern Leaf Blight': ('🔴 Fungal disease', 'Use resistant varieties. Apply fungicide. Rotate crops.'),
-    'Grape with Black Rot': ('🔴 Fungal disease', 'Apply myclobutanil or mancozeb. Remove mummified berries. Prune for airflow.'),
-    'Grape with Esca (Black Measles)': ('🔴 Fungal disease', 'No cure known. Remove infected wood. Avoid large pruning wounds.'),
-    'Grape with Isariopsis Leaf Spot': ('⚠️ Fungal disease', 'Apply copper fungicide. Remove infected plant debris.'),
-    'Orange with Citrus Greening': ('🔴 Bacterial disease (HLB)', 'No cure. Remove infected trees. Control Asian citrus psyllid vector.'),
-    'Peach with Bacterial Spot': ('⚠️ Bacterial disease', 'Apply copper bactericide. Use resistant varieties. Avoid overhead irrigation.'),
-    'Bell Pepper with Bacterial Spot': ('⚠️ Bacterial disease', 'Apply copper-based bactericide. Use disease-free seed. Rotate crops.'),
-    'Potato with Early Blight': ('⚠️ Fungal disease', 'Apply chlorothalonil or mancozeb. Ensure proper irrigation. Remove infected foliage.'),
-    'Potato with Late Blight': ('🔴 Fungal disease (P. infestans)', 'Apply metalaxyl or cymoxanil IMMEDIATELY. Destroy infected plants. This spreads rapidly!'),
-    'Squash with Powdery Mildew': ('⚠️ Fungal disease', 'Apply potassium bicarbonate or neem oil. Plant resistant varieties. Ensure spacing.'),
-    'Strawberry with Leaf Scorch': ('⚠️ Fungal disease', 'Apply captan fungicide. Remove infected leaves. Improve air circulation.'),
-    'Tomato with Bacterial Spot': ('⚠️ Bacterial disease', 'Apply copper-based bactericide. Use certified disease-free transplants.'),
-    'Tomato with Early Blight': ('⚠️ Fungal disease', 'Apply chlorothalonil. Remove lower leaves. Mulch around plants.'),
-    'Tomato with Late Blight': ('🔴 Fungal disease (P. infestans)', 'Apply metalaxyl IMMEDIATELY. Remove all infected material. Very contagious!'),
-    'Tomato with Leaf Mold': ('⚠️ Fungal disease', 'Improve ventilation. Reduce humidity. Apply chlorothalonil or copper.'),
-    'Tomato with Septoria Leaf Spot': ('⚠️ Fungal disease', 'Apply chlorothalonil. Remove infected leaves. Avoid overhead watering.'),
-    'Tomato with Spider Mites or Two-spotted Spider Mite': ('⚠️ Pest infestation', 'Apply miticide or neem oil. Increase humidity. Introduce predatory mites.'),
-    'Tomato with Target Spot': ('⚠️ Fungal disease', 'Apply azoxystrobin or chlorothalonil. Remove infected leaves.'),
-    'Tomato Yellow Leaf Curl Virus': ('🔴 Viral disease (TYLCV)', 'No cure. Remove infected plants. Control whitefly vectors. Use resistant varieties.'),
-    'Tomato Mosaic Virus': ('🔴 Viral disease', 'No cure. Remove infected plants. Disinfect tools. Control aphid vectors.'),
+    'Apple Scab': ('⚠️ Fungal', 'Apply fungicide at bud break (captan/myclobutanil). Remove infected leaves.'),
+    'Apple with Black Rot': ('🔴 Fungal', 'Prune infected wood. Apply copper fungicide. Improve air circulation.'),
+    'Cedar Apple Rust': ('⚠️ Fungal', 'Remove nearby cedar trees. Apply protective fungicide in spring.'),
+    'Cherry with Powdery Mildew': ('⚠️ Fungal', 'Apply sulfur spray. Avoid overhead watering.'),
+    'Corn (Maize) with Cercospora and Gray Leaf Spot': ('⚠️ Fungal', 'Use resistant varieties. Apply strobilurin fungicide.'),
+    'Corn (Maize) with Common Rust': ('⚠️ Fungal', 'Plant resistant hybrids. Apply fungicide at early signs.'),
+    'Corn (Maize) with Northern Leaf Blight': ('🔴 Fungal', 'Use resistant varieties. Apply fungicide. Rotate crops.'),
+    'Grape with Black Rot': ('🔴 Fungal', 'Apply myclobutanil. Remove mummified berries. Prune for airflow.'),
+    'Grape with Esca (Black Measles)': ('🔴 Fungal', 'No cure known. Remove infected wood. Avoid large pruning wounds.'),
+    'Grape with Isariopsis Leaf Spot': ('⚠️ Fungal', 'Apply copper fungicide. Remove infected debris.'),
+    'Orange with Citrus Greening': ('🔴 Bacterial (HLB)', 'No cure. Remove infected trees. Control psyllid vector.'),
+    'Peach with Bacterial Spot': ('⚠️ Bacterial', 'Apply copper bactericide. Use resistant varieties.'),
+    'Bell Pepper with Bacterial Spot': ('⚠️ Bacterial', 'Apply copper bactericide. Use disease-free seed. Rotate crops.'),
+    'Potato with Early Blight': ('⚠️ Fungal', 'Apply chlorothalonil. Ensure proper irrigation.'),
+    'Potato with Late Blight': ('🔴 Fungal (P. infestans)', 'Apply metalaxyl IMMEDIATELY. Destroy infected plants. Spreads rapidly!'),
+    'Squash with Powdery Mildew': ('⚠️ Fungal', 'Apply potassium bicarbonate or neem oil. Ensure plant spacing.'),
+    'Strawberry with Leaf Scorch': ('⚠️ Fungal', 'Apply captan fungicide. Remove infected leaves.'),
+    'Tomato with Bacterial Spot': ('⚠️ Bacterial', 'Apply copper bactericide. Use disease-free transplants.'),
+    'Tomato with Early Blight': ('⚠️ Fungal', 'Apply chlorothalonil. Remove lower leaves. Mulch around plants.'),
+    'Tomato with Late Blight': ('🔴 Fungal (P. infestans)', 'Apply metalaxyl IMMEDIATELY. Very contagious — remove all infected material!'),
+    'Tomato with Leaf Mold': ('⚠️ Fungal', 'Improve ventilation. Reduce humidity. Apply chlorothalonil.'),
+    'Tomato with Septoria Leaf Spot': ('⚠️ Fungal', 'Apply chlorothalonil. Remove infected leaves. Avoid overhead watering.'),
+    'Tomato with Spider Mites or Two-spotted Spider Mite': ('⚠️ Pest', 'Apply miticide or neem oil. Increase humidity.'),
+    'Tomato with Target Spot': ('⚠️ Fungal', 'Apply azoxystrobin. Remove infected leaves.'),
+    'Tomato Yellow Leaf Curl Virus': ('🔴 Viral (TYLCV)', 'No cure. Remove infected plants. Control whitefly vectors.'),
+    'Tomato Mosaic Virus': ('🔴 Viral', 'No cure. Remove infected plants. Disinfect tools.'),
 }
-HEALTHY_CROPS = {l for l in ID2LABEL.values() if 'Healthy' in l}
-
-# ── Transform (CPU-optimized, 2-view TTA) ────────────────────────────────────
-_TF_CENTER = T.Compose([
-    T.Resize(256),
-    T.CenterCrop(224),
-    T.ToTensor(),
-    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-])
-_TF_HFLIP = T.Compose([
-    T.Resize(256),
-    T.CenterCrop(224),
-    T.RandomHorizontalFlip(p=1.0),
-    T.ToTensor(),
-    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-])
-
-# ── Model cache ───────────────────────────────────────────────────────────────
-_models: dict = {}
 
-def _build_effnet(num_classes: int = 38) -> nn.Module:
+# ── Model loading ─────────────────────────────────────────────────────────────
+_model = None
+_load_status = "Not loaded"
+
+HF_BASE = "https://huggingface.co/noivan/cropdoc/resolve/main"
+MODEL_PREFIX = "models/cropdoc_efficientnet_v2/model.pt"
+N_PARTS = 16
+
+def _download_and_reconstruct():
+    """Download model parts from HF Hub and reconstruct."""
+    global _load_status
+    _load_status = "Downloading model parts..."
+    print(_load_status, flush=True)
+    
+    buf = io.BytesIO()
+    for i in range(N_PARTS):
+        part_url = f"{HF_BASE}/{MODEL_PREFIX}.part{i:04d}.json"
+        _load_status = f"Downloading part {i+1}/{N_PARTS}..."
+        print(_load_status, flush=True)
+        
+        for attempt in range(3):
+            try:
+                r = requests.get(part_url, timeout=60)
+                r.raise_for_status()
+                chunk_data = r.json()
+                buf.write(base64.b64decode(chunk_data["data"]))
+                break
+            except Exception as e:
+                if attempt == 2:
+                    raise RuntimeError(f"Failed to download part {i}: {e}")
+                time.sleep(2)
+    
+    buf.seek(0)
+    _load_status = "Loading model weights..."
+    print(_load_status, flush=True)
+    ckpt = torch.load(buf, map_location="cpu", weights_only=False)
+    return ckpt
+
+def _build_model(ckpt):
+    nc = ckpt.get("num_classes", 38)
     m = efficientnet_v2_s(weights=None)
-    m.classifier[1] = nn.Linear(m.classifier[1].in_features, num_classes)
+    m.classifier[1] = nn.Linear(m.classifier[1].in_features, nc)
+    if "model_state_dict" in ckpt:
+        m.load_state_dict(ckpt["model_state_dict"])
+    else:
+        m.load_state_dict(ckpt)
+    m.eval()
     return m
 
-def _load_model(key: str, filename: str) -> nn.Module | None:
-    if key in _models:
-        return _models[key]
+def ensure_model():
+    global _model, _load_status
+    if _model is not None:
+        return True
     try:
-        path = hf_hub_download(repo_id="noivan/cropdoc", filename=filename)
-        ckpt = torch.load(path, map_location="cpu", weights_only=False)
-        if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
-            state = ckpt["model_state_dict"]
-            nc = ckpt.get("num_classes", 38)
-        else:
-            state = ckpt
-            nc = 38
-        model = _build_effnet(nc)
-        model.load_state_dict(state)
-        model.eval()
-        _models[key] = model
-        print(f"✅ Loaded {key}")
-        return model
+        ckpt = _download_and_reconstruct()
+        _model = _build_model(ckpt)
+        _load_status = "✅ Model loaded (EfficientNetV2-S, 99.9% PlantVillage)"
+        print(_load_status, flush=True)
+        return True
     except Exception as e:
-        print(f"⚠️ Failed to load {key}: {e}")
-        return None
-
-def _predict_tta(img: Image.Image, model: nn.Module) -> np.ndarray:
-    """2-view TTA (Center + HFlip), batched for CPU efficiency."""
-    t1 = _TF_CENTER(img)
-    t2 = _TF_HFLIP(img)
-    batch = torch.stack([t1, t2])  # (2, 3, 224, 224)
-    with torch.no_grad():
-        logits = model(batch)
-        probs = torch.softmax(logits, dim=1).mean(dim=0)  # average over TTA
-    return probs.numpy()
-
-def _ensemble_predict(img: Image.Image) -> np.ndarray | None:
-    """FT v3+v4 ensemble (v3=0.40, v4=0.60). v2 optional if loaded."""
-    ft_v2 = _models.get("ft_v2")
-    ft_v3 = _models.get("ft_v3")
-    ft_v4 = _models.get("ft_v4")
-
-    if ft_v3 is None and ft_v4 is None:
-        return None
-
-    probs = None
-    if ft_v2 is not None and ft_v3 is not None and ft_v4 is not None:
-        p2 = _predict_tta(img, ft_v2)
-        p3 = _predict_tta(img, ft_v3)
-        p4 = _predict_tta(img, ft_v4)
-        probs = 0.10 * p2 + 0.50 * p3 + 0.40 * p4
-    elif ft_v3 is not None and ft_v4 is not None:
-        p3 = _predict_tta(img, ft_v3)
-        p4 = _predict_tta(img, ft_v4)
-        probs = 0.40 * p3 + 0.60 * p4
-    elif ft_v3 is not None:
-        probs = _predict_tta(img, ft_v3)
-    elif ft_v4 is not None:
-        probs = _predict_tta(img, ft_v4)
-
-    if probs is None:
-        return None
-
-    # Apply class calibration (v40)
-    probs = probs.copy()
-    for idx, mult in CALIBRATION.items():
-        probs[idx] *= mult
-
-    return probs
-
-# ── Lazy model loading on first inference ─────────────────────────────────────
-_initialized = False
-
-def _ensure_models():
-    global _initialized
-    if _initialized:
-        return
-    _initialized = True
-    # Primary: ft_v3 + ft_v4 (most impactful)
-    _load_model("ft_v3", "models/effnet_ft_v3.pt")
-    _load_model("ft_v4", "models/effnet_ft_v4.pt")
-    # Optional: ft_v2 (small weight 0.10)
-    _load_model("ft_v2", "models/effnet_ft_plantdoc.pt")
-
-# ── Main inference function ───────────────────────────────────────────────────
-def diagnose(image: Image.Image) -> tuple[str, str, str, dict]:
-    """
-    Returns: (label, severity, recommendation, confidences_dict)
-    """
-    if image is None:
-        return "No image", "", "Please upload an image.", {}
-
+        _load_status = f"❌ Load failed: {e}"
+        print(_load_status, flush=True)
+        return False
+
+# ── Transforms ───────────────────────────────────────────────────────────────
+_TF_C = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor(),
+    T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
+_TF_H = T.Compose([T.Resize(256), T.CenterCrop(224),
+    T.RandomHorizontalFlip(1.0), T.ToTensor(),
+    T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
+
+def predict(img: Image.Image):
     t0 = time.time()
-    _ensure_models()
-
-    img = image.convert("RGB")
+    if not ensure_model():
+        return "Model failed to load", "❌", _load_status, {}
+    
+    img = img.convert("RGB")
     if max(img.size) > 1024:
         img.thumbnail((1024, 1024), Image.LANCZOS)
-
-    probs = _ensemble_predict(img)
-    if probs is None:
-        return "Model loading failed", "❌", "Please try again.", {}
-
-    # Top-5 results
+    
+    t1, t2 = _TF_C(img), _TF_H(img)
+    batch = torch.stack([t1, t2])
+    with torch.no_grad():
+        probs = torch.softmax(_model(batch), dim=1).mean(0).numpy()
+    
     top5 = sorted(enumerate(probs), key=lambda x: -x[1])[:5]
-    top1_idx, top1_score = top5[0]
-    top1_label = ID2LABEL.get(top1_idx, "Unknown")
-
+    idx, score = top5[0]
+    label = ID2LABEL.get(idx, "Unknown")
     elapsed = time.time() - t0
-
-    # Determine severity / recommendation
-    if top1_label in HEALTHY_CROPS:
-        severity = "✅ Healthy"
+    
+    if label in HEALTHY:
+        sev = "✅ Healthy"
         rec = "Your plant looks healthy! Continue regular care and monitoring."
     else:
-        info = DISEASE_INFO.get(top1_label)
-        if info:
-            severity, rec = info
-        else:
-            severity = "⚠️ Disease detected"
-            rec = "Consult a local agricultural extension service."
-
-    # Build confidence dict for display
-    conf_dict = {ID2LABEL[i]: float(f"{s:.4f}") for i, s in top5}
-
-    # Format label with confidence
-    label_out = f"{top1_label} ({top1_score:.1%})"
-    rec_out = f"{rec}\n\n⏱ Inference: {elapsed:.2f}s (CPU)"
-
-    return label_out, severity, rec_out, conf_dict
+        info = DISEASE_INFO.get(label, ("⚠️ Disease detected", "Consult a local agricultural extension service."))
+        sev, rec = info
+    
+    conf = {ID2LABEL[i]: float(f"{s:.4f}") for i, s in top5}
+    label_out = f"{label} ({score:.1%})"
+    rec_out = f"{rec}\n\n⏱ {elapsed:.2f}s · CPU · EfficientNetV2-S"
+    return label_out, sev, rec_out, conf
 
 # ── Gradio UI ─────────────────────────────────────────────────────────────────
 CSS = """
 .gradio-container { max-width: 900px !important; margin: auto; }
-.title { text-align: center; color: #0d5c32; }
-.subtitle { text-align: center; color: #2a6040; margin-bottom: 1em; }
-#result-label { font-size: 1.3em; font-weight: bold; }
-.healthy { color: #007830; }
-.disease { color: #cc4400; }
+h1 { text-align: center; color: #0d5c32; }
+.subtitle { text-align: center; color: #2a6040; margin-bottom: 1em; font-size: 14px; }
 footer { display: none !important; }
 """
 
-EXAMPLES = [
-    # These will be populated if example images are added to the repo
-]
-
-def run_diagnose(img):
-    if img is None:
-        return "Upload an image first", "", "", {}
-    label, severity, rec, conf = diagnose(img)
-    return label, severity, rec, conf
-
 with gr.Blocks(css=CSS, title="🌿 CropDoc — Plant Disease AI") as demo:
     gr.HTML("""
-    <div class="title">
-      <h1>🌿 CropDoc</h1>
-      <p class="subtitle">Plant Disease Diagnosis AI · 38 Disease Classes · EfficientNetV2-S Ensemble</p>
-    </div>
+    <h1>🌿 CropDoc</h1>
+    <p class="subtitle">Plant Disease Diagnosis AI · 38 Classes · EfficientNetV2-S (99.9% accuracy)</p>
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            img_input = gr.Image(type="pil", label="📷 Upload plant photo")
+            img_in = gr.Image(type="pil", label="📷 Upload plant photo")
             btn = gr.Button("🔍 Diagnose", variant="primary", size="lg")
-            gr.HTML("""
-            <div style="font-size:12px;color:#555;margin-top:8px;">
-            Supported crops: Tomato, Potato, Pepper, Apple, Grape, Corn, Peach, Strawberry, Squash, Orange, Cherry, Raspberry, Soybean, Blueberry<br>
-            <b>Note:</b> Running on CPU — first inference ~5s (model loading), subsequent ~2s
-            </div>
-            """)
+            gr.HTML("""<div style="font-size:12px;color:#777;margin-top:8px;">
+            <b>Supported crops:</b> Tomato, Potato, Pepper, Apple, Grape, Corn, Peach, Strawberry, Squash, Orange, Cherry, Raspberry, Soybean, Blueberry<br>
+            <b>Note:</b> First run downloads model (~80MB). Subsequent inferences are fast (~1–2s CPU).
+            </div>""")
 
         with gr.Column(scale=1):
-            out_label = gr.Textbox(label="🏷️ Diagnosis", elem_id="result-label")
-            out_severity = gr.Textbox(label="⚡ Severity")
-            out_rec = gr.Textbox(label="💊 Recommendation", lines=4)
-            out_conf = gr.JSON(label="📊 Top-5 Confidence")
-
-    btn.click(
-        fn=run_diagnose,
-        inputs=[img_input],
-        outputs=[out_label, out_severity, out_rec, out_conf],
-    )
-    img_input.change(
-        fn=run_diagnose,
-        inputs=[img_input],
-        outputs=[out_label, out_severity, out_rec, out_conf],
-    )
-
-    gr.HTML("""
-    <hr style="margin-top:2em;">
-    <div style="text-align:center;font-size:12px;color:#666;">
-      <b>CropDoc</b> · EfficientNetV2-S fine-tuned on PlantDoc (61.6% accuracy) ·
-      <a href="https://github.com/noivan0/cropdoc" target="_blank">GitHub</a> ·
-      <a href="https://huggingface.co/noivan/cropdoc" target="_blank">Model Weights</a>
-    </div>
-    """)
+            out_label = gr.Textbox(label="🏷️ Diagnosis")
+            out_sev   = gr.Textbox(label="⚡ Severity")
+            out_rec   = gr.Textbox(label="💊 Recommendation", lines=4)
+            out_conf  = gr.JSON(label="📊 Top-5 Confidence")
+
+    def run(img):
+        if img is None:
+            return "Upload an image first", "", "", {}
+        return predict(img)
+
+    btn.click(run, [img_in], [out_label, out_sev, out_rec, out_conf])
+    img_in.change(run, [img_in], [out_label, out_sev, out_rec, out_conf])
+
+    gr.HTML("""<hr style="margin-top:2em;">
+    <div style="text-align:center;font-size:12px;color:#888;">
+    CropDoc · EfficientNetV2-S · PlantVillage 99.9% · PlantDoc 61.6% ·
+    <a href="https://github.com/noivan0/cropdoc">GitHub</a> ·
+    <a href="https://huggingface.co/noivan/cropdoc">Model</a>
+    </div>""")
 
 if __name__ == "__main__":
     demo.launch()