File size: 3,186 Bytes
88b9d97 d36fd45 88b9d97 8ee30a6 88b9d97 8ee30a6 88b9d97 8ee30a6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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")
gr.Markdown("### A semantic segmentation model for counting wheat heads in field images.")
gr.Markdown("- Designed for yield estimation, flowering time detection, and field maturity assessment")
gr.Markdown("- Best results with overhead imagery under diffuse lighting")
gr.Markdown("- Performance may vary with harsh lighting or heavily overlapping 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.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]
)
gr.Markdown("<div style='text-align: center'>Powered by <strong><a href='https://huggingface.co/chmcbs/HeadCount' target='_blank'>chmcbs/HeadCount</a></strong></div>")
if __name__ == "__main__":
demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
|