Add model weights, inference code, and dependencies

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}")

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a44b1504d0ce10a601cda4adf11bbc967d87f42bb2b2622fa922b5667bcaf17
+size 320679895

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.0.0
+torchvision>=0.15.0
+segmentation-models-pytorch>=0.3.3
+Pillow>=9.0.0
+numpy>=1.24.0
+scipy>=1.10.0
+scikit-image>=0.20.0