Add Gradio application and inference code

app.py

@@ -0,0 +1,81 @@
+import gradio as gr
+import os
+from inference import GWFSSModel
+from PIL import Image
+import numpy as np
+from scipy import ndimage
+from skimage.feature import peak_local_max
+from huggingface_hub import hf_hub_download
+
+# Download model from Hugging Face
+print("Downloading model from Hugging Face...")
+MODEL_PATH = hf_hub_download(repo_id="chmcbs/HeadCount", filename="model.pth")
+print("✓ Model downloaded successfully")
+
+# Load model
+print("Loading model...")
+model = GWFSSModel(MODEL_PATH)
+print("✓ Model loaded successfully")
+
+def process_image(image):
+    if image is None:
+        return "", None
+    
+    try:
+        predictions = model.predict(image)
+        num_heads = model.count_heads(predictions)
+        
+        # Visualise detected peaks
+        head_mask = (predictions == 3).astype(np.uint8)
+        distance = ndimage.distance_transform_edt(head_mask)
+        coords = peak_local_max(distance, min_distance=15, labels=head_mask)
+        
+        # Create overlay with peak markers
+        overlay = model.overlay_mask(image, predictions, alpha=0.5, heads_only=True)
+        overlay_np = np.array(overlay)
+        for y, x in coords:
+            # Draw a small red circle at each detected peak
+            overlay_np[max(0,y-3):y+4, max(0,x-3):x+4] = [255, 0, 0]
+        
+        overlay = Image.fromarray(overlay_np)
+        count_message = f"### 🌾 {num_heads} heads detected!"
+        return count_message, overlay
+    except Exception as e:
+        return f"Error: {str(e)}", None
+
+# Get example images
+example_images = []
+if os.path.exists("examples"):
+    example_files = sorted([f for f in os.listdir("examples") 
+                           if f.endswith(('.jpg', '.jpeg', '.png'))])[:5]
+    example_images = [os.path.join("examples", f) for f in example_files]
+
+# Create Gradio interface
+with gr.Blocks(title="HeadCount") as demo:
+    gr.Markdown("# 🌾 HeadCount: Automated Wheat Head Counter")
+    gr.Markdown("Upload an image to automatically detect and count wheat heads.")
+    
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+        
+        with gr.Column():
+            overlay_output = gr.Image(label="Segmentation Overlay")
+            generate_btn = gr.Button("Generate", variant="primary")
+    
+    with gr.Row():
+        with gr.Column():
+            if example_images:
+                gr.Markdown("### Example Images")
+                gr.Examples(examples=example_images, inputs=image_input)
+        
+        with gr.Column():
+            head_count_output = gr.Markdown(value="")
+    
+    generate_btn.click(
+        fn=process_image,
+        inputs=image_input,
+        outputs=[head_count_output, overlay_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)

inference.py

@@ -0,0 +1,173 @@
+"""
+Inference module for counting wheat heads in field images using a DeepLabV3+ semantic
+segmentation model trained on the GWFSS dataset.
+
+The model performs multi-class segmentation (Background, Leaf, Stem, Head) to accurately
+distinguish wheat heads from other plant organs, then uses connected component analysis
+to count individual heads.
+"""
+
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+import segmentation_models_pytorch as smp
+from scipy import ndimage
+from skimage.feature import peak_local_max
+
+# ImageNet normalisation constants
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+# Mask colours for visualization
+MASK_COLORS = [
+    (0, 0, 0),          # Background: black
+    (214, 255, 50),     # Leaf: yellow-green
+    (50, 132, 255),     # Stem: blue
+    (50, 255, 132),     # Head: cyan-green
+]
+
+class GWFSSModel:
+    def __init__(self, model_path, device=None):
+        if device is None:
+            if torch.cuda.is_available():
+                self.device = torch.device("cuda")
+            elif torch.backends.mps.is_available():
+                self.device = torch.device("mps")
+            else:
+                self.device = torch.device("cpu")
+        else:
+            self.device = device
+        
+        # Load model architecture
+        self.model = smp.DeepLabV3Plus(
+            encoder_name="resnet50",
+            encoder_weights=None,
+            in_channels=3,
+            classes=4,
+        )
+        
+        # Load trained weights
+        checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        
+        # Image preprocessing
+        self.transform = transforms.Compose([
+            transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+        ])
+    
+    def preprocess_image(self, image):
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        image_tensor = self.transform(image).unsqueeze(0)
+        return image_tensor.to(self.device)
+
+    def predict(self, image):
+        if isinstance(image, str):
+            image = Image.open(image)
+        
+        image_tensor = self.preprocess_image(image)
+        
+        with torch.no_grad():
+            logits = self.model(image_tensor)
+        
+        predictions = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy()
+        return predictions
+
+    def count_heads(self, predictions, min_distance=15):
+        head_mask = (predictions == 3).astype(np.uint8)
+        
+        if head_mask.sum() == 0:
+            return 0
+        
+        # Compute distance transform
+        distance = ndimage.distance_transform_edt(head_mask)
+        
+        # Find local peaks (head centers)
+        coords = peak_local_max(distance, min_distance=min_distance, labels=head_mask)
+        
+        # Count the peaks
+        num_heads = len(coords)
+        
+        return num_heads
+
+    def create_colored_mask(self, predictions):
+        h, w = predictions.shape
+        mask_rgb = np.zeros((h, w, 3), dtype=np.uint8)
+        
+        for class_id, color in enumerate(MASK_COLORS):
+            mask_rgb[predictions == class_id] = color
+        
+        return Image.fromarray(mask_rgb)
+
+    def overlay_mask(self, image, predictions, alpha=0.5, heads_only=True):
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
+        if image.size != (512, 512):
+            image = image.resize((512, 512), Image.Resampling.BILINEAR)
+        
+        # Create mask
+        h, w = predictions.shape
+        mask_rgb = np.zeros((h, w, 3), dtype=np.uint8)
+        
+        if heads_only:
+            # Only highlight heads
+            mask_rgb[predictions == 3] = (50, 255, 132)
+        else:
+            # Show all classes
+            for class_id, color in enumerate(MASK_COLORS):
+                mask_rgb[predictions == class_id] = color
+        
+        mask_img = Image.fromarray(mask_rgb)
+        overlay = Image.blend(image.convert('RGB'), mask_img, alpha)
+        return overlay
+
+    def predict_and_overlay(self, image, alpha=0.5, heads_only=True):
+        predictions = self.predict(image)
+        overlay = self.overlay_mask(image, predictions, alpha=alpha, heads_only=heads_only)
+        return overlay
+
+if __name__ == "__main__":
+    import sys
+    
+    if len(sys.argv) < 2:
+        print("Usage: python inference.py <image_path> [model_path]")
+        sys.exit(1)
+
+    image_path = sys.argv[1]
+    model_path = sys.argv[2] if len(sys.argv) > 2 else "cache/02_dice_stem.pth"
+    
+    print(f"Loading model from {model_path}...")
+    model = GWFSSModel(model_path)
+    
+    print(f"Processing image: {image_path}")
+    image = Image.open(image_path)
+    predictions = model.predict(image)
+    
+    # Count heads
+    num_heads = model.count_heads(predictions)
+    print(f"\n🌾 {num_heads} heads detected!")
+    
+    # Create visualisations
+    print("\nGenerating visualisations...")
+    overlay_heads = model.overlay_mask(image, predictions, alpha=0.5, heads_only=True)
+    overlay_all = model.overlay_mask(image, predictions, alpha=0.5, heads_only=False)
+    
+    # Save outputs
+    output_heads = image_path.rsplit('.', 1)[0] + '_heads_only.png'
+    output_all = image_path.rsplit('.', 1)[0] + '_all_classes.png'
+    
+    overlay_heads.save(output_heads)
+    overlay_all.save(output_all)
+    
+    print(f"✓ Saved head overlay to: {output_heads}")
+    print(f"✓ Saved full segmentation to: {output_all}")