update app: SAM2 vit_b + SAM3 integration

README.md

@@ -1,6 +1,6 @@
 ---
-title: CLR Severity
-emoji: 😻
+title: CLR Severity Estimator
+emoji: ☕
 colorFrom: blue
 colorTo: green
 sdk: gradio
@@ -8,7 +8,17 @@ sdk_version: 6.9.0
 app_file: app.py
 pinned: false
 license: mit
-short_description: Coffee Leaf Rust leaf and severity segmentation framework.
+short_description: Coffee Leaf Rust pipeline using YOLOv8, SAM2, and SAM3.
 ---
 
+# ☕ Coffee Leaf Rust (CLR) Severity Estimator
+
+This framework processes coffee leaf images to accurately estimate rust severity using a 3-step deep learning pipeline:
+
+1. **Leaf Detection**: **YOLOv8** locates and extracts bounding boxes for all coffee leaves in the image.
+2. **Instance Segmentation**: **SAM2 (Segment Anything Model)** takes the bounding boxes to create pixel-perfect black-background cutouts of each leaf.
+3. **Rust Segmentation**: **SAM3** uses a text prompt ("yellow spot") to find and segment the rust lesions on each leaf cutout.
+
+The Gradio interface presents a detailed summary table and visualizations for each individual leaf in the image.
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -4,194 +4,246 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
+import urllib.request
 from ultralytics import YOLO
 
+# Try importing SAM2
+try:
+    from segment_anything import sam_model_registry, SamPredictor
+    SAM2_AVAILABLE = True
+except ImportError:
+    SAM2_AVAILABLE = False
+    print("SAM2 not available")
+
+# Try importing SAM3
+try:
+    from sam3.model_builder import build_sam3_image_model
+    from sam3.model.sam3_image_processor import Sam3Processor
+    SAM3_AVAILABLE = True
+except ImportError:
+    SAM3_AVAILABLE = False
+    print("SAM3 not available")
+
 ############################################
 # Configuration
 ############################################
 
-YOLO_MODEL_PATH = "clr_YOLOV8.pt"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print("Running on:", DEVICE)
+
+YOLO_MODEL_PATH = "clr_YOLOV8.pt"
+
+# SAM2 vit_b (lighter)
+SAM2_MODEL_TYPE = "vit_b"
+SAM2_CHECKPOINT_PATH = "sam_vit_b_01ec64.pth"
 
 ############################################
-# Load YOLO Model
+# Download SAM2 if needed
 ############################################
 
-print("Loading YOLO model...")
+if SAM2_AVAILABLE:
+    if not os.path.exists(SAM2_CHECKPOINT_PATH):
+        print("Downloading SAM2 vit_b checkpoint...")
+        urllib.request.urlretrieve(
+            "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+            SAM2_CHECKPOINT_PATH
+        )
 
+############################################
+# Load Models
+############################################
+
+print("Loading models...")
+
+# YOLO
+yolo_model = None
 try:
-    yolo_model = YOLO(YOLO_MODEL_PATH)
-    print("YOLO model loaded.")
+    if os.path.exists(YOLO_MODEL_PATH):
+        yolo_model = YOLO(YOLO_MODEL_PATH)
+        print("YOLO loaded")
+    else:
+        print("YOLO model not found")
 except Exception as e:
-    yolo_model = None
-    print("YOLO loading error:", e)
+    print("YOLO error:", e)
+
+# SAM2
+sam2_predictor = None
+if SAM2_AVAILABLE:
+    try:
+        sam2 = sam_model_registry[SAM2_MODEL_TYPE](
+            checkpoint=SAM2_CHECKPOINT_PATH
+        )
+        sam2.to(DEVICE)
+        sam2_predictor = SamPredictor(sam2)
+        print("SAM2 loaded")
+    except Exception as e:
+        print("SAM2 error:", e)
+
+# SAM3 (official, no checkpoint)
+sam3_processor = None
+if SAM3_AVAILABLE:
+    try:
+        sam3_model = build_sam3_image_model(device=DEVICE)
+        sam3_processor = Sam3Processor(sam3_model)
+        print("SAM3 loaded")
+    except Exception as e:
+        print("SAM3 error:", e)
 
 ############################################
 # Helper Functions
 ############################################
 
-def segment_rust_simple(leaf_img):
-    """
-    Simple rust segmentation using HSV color threshold.
-    Works as fallback when SAM is unavailable.
-    """
-
+def fallback_segment_rust(leaf_img):
     hsv = cv2.cvtColor(leaf_img, cv2.COLOR_BGR2HSV)
-
-    # Rust-like colors
     lower = np.array([10, 80, 80])
     upper = np.array([35, 255, 255])
-
     mask = cv2.inRange(hsv, lower, upper)
-
     kernel = np.ones((3,3), np.uint8)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    return cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+
+def extract_leaf_sam2(image_rgb, box):
+    if not sam2_predictor:
+        return np.ones((image_rgb.shape[0], image_rgb.shape[1]), dtype=np.uint8) * 255
+
+    sam2_predictor.set_image(image_rgb)
+    masks, _, _ = sam2_predictor.predict(
+        box=np.array(box),
+        multimask_output=False
+    )
+    return (masks[0] * 255).astype(np.uint8)
+
+def segment_lesions_sam3(image_rgb):
+    if not sam3_processor:
+        return fallback_segment_rust(cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR))
+
+    try:
+        pil_img = Image.fromarray(image_rgb)
+        state = sam3_processor.set_image(pil_img)
+        output = sam3_processor.set_text_prompt(
+            state=state,
+            prompt="yellow spot"
+        )
 
-    return mask
+        masks = output.get("masks", None)
 
+        if masks is None or len(masks) == 0:
+            return np.zeros((image_rgb.shape[0], image_rgb.shape[1]), dtype=np.uint8)
 
-def calculate_leaf_area(leaf_img):
-    """
-    Estimate leaf pixels via threshold.
-    """
-    gray = cv2.cvtColor(leaf_img, cv2.COLOR_BGR2GRAY)
-    _, mask = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
+        combined = None
+        for m in masks:
+            m_np = m.detach().cpu().numpy()
+            m_np = np.squeeze(m_np)
+            m_np = (m_np > 0).astype(np.uint8)
 
-    return mask
+            if combined is None:
+                combined = m_np
+            else:
+                combined = np.maximum(combined, m_np)
 
+        return combined * 255
+
+    except Exception as e:
+        print("SAM3 error:", e)
+        return fallback_segment_rust(cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR))
 
 ############################################
-# Main Processing Function
+# Main Function
 ############################################
 
 def process_coffee_leaf(image):
-
     if image is None:
-        return None, "Upload an image."
+        return None, None, [["Upload image", "-", "-"]]
 
     if yolo_model is None:
-        return image, "YOLO model not loaded."
+        return image, None, [["Error", "YOLO not loaded", "-"]]
 
     image_np = np.array(image)
     image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    h, w = image_cv.shape[:2]
 
     results = yolo_model(image_cv, verbose=False)
-
     boxes = results[0].boxes.xyxy.cpu().numpy()
 
     if len(boxes) == 0:
-        return image_np, "No leaves detected."
+        return image_np, None, [["No leaves detected", "-", "-"]]
 
     annotated = image_np.copy()
-
-    severities = []
-
-    h, w = image_cv.shape[:2]
+    gallery = []
+    table = []
 
     for i, box in enumerate(boxes):
-
         x1, y1, x2, y2 = box.astype(int)
 
-        x1, x2 = max(0, x1), min(w, x2)
-        y1, y2 = max(0, y1), min(h, y2)
-
-        leaf_crop = image_cv[y1:y2, x1:x2]
-
+        leaf_crop = image_np[y1:y2, x1:x2]
         if leaf_crop.size == 0:
             continue
 
-        ################################
-        # Leaf mask
-        ################################
+        # SAM2 leaf mask
+        leaf_mask_full = extract_leaf_sam2(image_np, box)
+        leaf_mask = leaf_mask_full[y1:y2, x1:x2]
 
-        leaf_mask = calculate_leaf_area(leaf_crop)
         leaf_pixels = cv2.countNonZero(leaf_mask)
 
-        ################################
-        # Rust segmentation
-        ################################
-
-        rust_mask = segment_rust_simple(leaf_crop)
-        rust_pixels = cv2.countNonZero(rust_mask)
-
-        ################################
-        # Severity calculation
-        ################################
-
-        severity = 0
-
-        if leaf_pixels > 0:
-            severity = (rust_pixels / leaf_pixels) * 100
+        # Cutout
+        cutout = np.zeros_like(leaf_crop)
+        cutout[leaf_mask > 0] = leaf_crop[leaf_mask > 0]
 
-        severities.append(f"Leaf {i+1}: {severity:.2f}%")
+        # SAM3 lesions
+        rust_mask = segment_lesions_sam3(cutout)
+        rust_mask = cv2.bitwise_and(rust_mask, leaf_mask)
 
-        ################################
-        # Visualization
-        ################################
+        rust_pixels = cv2.countNonZero(rust_mask)
 
-        # draw leaf bbox
-        cv2.rectangle(annotated,(x1,y1),(x2,y2),(0,255,0),2)
+        severity = (rust_pixels / leaf_pixels) * 100 if leaf_pixels > 0 else 0
 
-        cv2.putText(
-            annotated,
-            f"{severity:.1f}%",
-            (x1,y1-5),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.6,
-            (0,255,0),
-            2
-        )
+        # Draw bbox
+        cv2.rectangle(annotated, (x1,y1), (x2,y2), (0,255,0), 2)
+        cv2.putText(annotated, f"{severity:.1f}%", (x1,y1-5),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
 
-        # resize rust mask to image coords
-        rust_mask = cv2.resize(
-            rust_mask,
-            (x2-x1, y2-y1),
-            interpolation=cv2.INTER_NEAREST
-        )
+        # Overlay
+        overlay = cutout.copy()
+        overlay[rust_mask > 0] = [128, 0, 128]
 
-        overlay = np.zeros_like(annotated)
+        gallery.append((cutout, f"Leaf {i+1}"))
+        gallery.append((overlay, f"{severity:.1f}%"))
 
-        overlay[y1:y2, x1:x2][rust_mask > 0] = [255,0,0]
+        table.append([str(i+1), f"{severity:.1f}%", f"{100-severity:.1f}%"])
 
-        alpha = 0.4
-        mask_indices = overlay[:,:,0] > 0
+    return annotated, gallery, table
 
-        annotated[mask_indices] = (
-            annotated[mask_indices]*(1-alpha) +
-            overlay[mask_indices]*alpha
-        ).astype(np.uint8)
+############################################
+# UI
+############################################
 
-    report = f"Detected {len(boxes)} leaves\n\n"
-    report += "\n".join(severities)
+with gr.Blocks() as demo:
+    gr.Markdown("# ☕ Coffee Leaf Rust Severity Estimator")
 
-    return annotated, report
+    image_input = gr.Image(type="pil")
+    submit = gr.Button("Run")
+    clear = gr.Button("Clear")
 
+    output_img = gr.Image()
+    gallery = gr.Gallery(columns=2)
+    table = gr.Dataframe(headers=["Leaf", "Severity", "Healthy"])
 
-############################################
-# Gradio Interface
-############################################
+    submit.click(
+        process_coffee_leaf,
+        inputs=image_input,
+        outputs=[output_img, gallery, table]
+    )
 
-demo = gr.Interface(
-    fn=process_coffee_leaf,
-    inputs=gr.Image(type="pil", label="Upload Coffee Leaf Image"),
-    outputs=[
-        gr.Image(label="Analyzed Image"),
-        gr.Textbox(label="Severity Report")
-    ],
-    title="☕ Coffee Leaf Rust Severity Estimator",
-    description="""
-Upload a coffee leaf image.  
-The system detects leaves using YOLOv8 and estimates rust severity by segmenting rust-colored lesions.
-""",
-)
+    def clear_all():
+        return None, None, None, None
 
+    clear.click(
+        clear_all,
+        outputs=[image_input, output_img, gallery, table]
+    )
 
 ############################################
 # Launch
 ############################################
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+    demo.launch()