try: import torch import torchvision except ImportError: import subprocess print("Attempting to install missing packages...") subprocess.check_call(["pip", "install", "torch", "torchvision"]) import torch import torchvision import gradio as gr import os import numpy as np import torch import torch.nn as nn from torchvision import transforms import requests import os from PIL import Image from collections import OrderedDict from torchvision import models import torch.nn.functional as F import matplotlib.pyplot as plt import cv2 import io # Import CSS and URL File css_file_path = os.path.join(os.path.dirname(__file__), "ui.css") with open(css_file_path,"r") as f: custom_css = f.read() # HTML Design html_welcome_page = """

Welcome to RemoveWeed Weed Detection System

Project Aim: This system is designed to optimize rice planting schedules with broad-leaved weed detection using machine learning.

""" html_system_page ="""

RemoveWeed System Overview

This system is designed to help farmers detect broad-leaved weeds in rice fields using machine learning techniques. The aim is to optimize rice planting schedules and improve crop yield.

""" html_project_description = """

- 🌿 About Project 🌿 -

Resource: Statista (2024) - Agricultural consumption of herbicides worldwide from 1990 to 2022 (in 1,000 metric tons)

Herbicide Use Soars: A Shocking Yearly Increase!

Statista (2024) revealed that global herbicide consumption has reached 1.94 million metric tons. To control dock weed in farming fields, the application of herbicides can cause delays in rice planting schedules ranging from 7 to 30 days.

Why Choose RemoveWeed?

RemoveWeed is a system designed to detect broad-leaved dock weed in paddy fields. It uses object detection like Single Shot Detection (SSD) model, along with instance segmentation models like U-Net and Fully Convolutional Neural Network (FCNN, to predict the presence of dock weed.

Broad-leaved Dock Weed in Paddy Field

Potential Benefits

Cost Savings 💰
Reduce Labor and Manual Monitoring Cost 💹
Increase Profitability by Rice Planting Scheduling Advice 📈
Provide Sustainable Practices in Agriculture 🧑‍🌾
Reduce Herbicide Pollution ☢️

""" html_author_review_page = """

- Project Owner Introduction -

Whitney Lim Wan Yee is a student at Asia Pacific University (APU), pursuing Year 3 Computer Science specialization in Data Analytics. She is passionate about machine learning and its applications in agriculture.

""" js_func = """ function refresh() { const url = new URL(window.location); if (url.searchParams.get('__theme') !== 'light') { url.searchParams.set('__theme', 'light'); window.location.href = url.href; } } """ def choose_model(choice): if choice == "Instance Segmentation Model (U-Net)": return "You have selected U-Net" else: return "Invalid selection" # Gradio Interface def gradio_interface(selected_model, uploaded_image): # This will call the predict function and display the results return predict(selected_model, uploaded_image) with gr.Blocks(css=custom_css,js=js_func) as demo: # State to track current page page = gr.State(value="welcome") # Welcome page container with gr.Group(visible=True, elem_classes="gradio-container") as welcome_page: gr.HTML(html_welcome_page) # Insert HTML structure start_trial_button = gr.Button("Start Trial", variant="primary", elem_classes="trial-button") # System description page container (initially hidden) with gr.Group(visible=False) as system_page: gr.HTML(html_system_page) tabs = gr.Tabs() with tabs: with gr.TabItem("Project Description"): tab_state = gr.State(value=0) gr.HTML(html_project_description) with gr.TabItem("Model Playground"): gr.Markdown(""" ### Model Playground: This section allows users to interact with the model and test its capabilities. """) # Model selection radio buttons radio = gr.Radio(choices=["Instance Segmentation Model (U-Net)"], label="Select Model", elem_classes="model-selection") # Output for model selection result output = gr.Textbox(label="Model Selection Result", elem_classes="model-selection-output") # Trigger the choose_model function when a model is selected radio.change(fn=choose_model, inputs=radio, outputs=output) # Image input and upload button img_input = gr.Image(type="numpy", label="Upload Image", elem_classes="image-input") upload_image_button = gr.Button("Start Prediction", variant="primary", elem_classes="upload-button") # Predicted output img_output = gr.Image(label="Predicted Image", elem_classes="image-output") # Predict and show output when image is uploaded upload_image_button.click(fn=gradio_interface, inputs=[radio, img_input], outputs=img_output) with gr.TabItem("Open Source API Link"): gr.Markdown(""" ### Open Source API Link: This section provides access to the open-source API for the weed detection model. """) gr.Markdown("### API Documentation:") with gr.TabItem("Contact and Review"): gr.HTML(html_author_review_page) back_button = gr.Button("Back", variant="secondary",elem_classes="back-button") # Navigation functions def go_to_system_page(): print("Going to system page") return gr.update(visible=False), gr.update(visible=True) def go_to_welcome_page(): print("Going to welcome page") return gr.update(visible=True), gr.update(visible=False) def process_image(uploaded_image): # If the image is passed as a numpy array, convert it to a PIL image if isinstance(uploaded_image, np.ndarray): image = Image.fromarray(uploaded_image) elif isinstance(uploaded_image, Image.Image): image = uploaded_image else: raise ValueError("Uploaded image must be either a numpy array or a PIL Image.") # Define the necessary transformations transform = transforms.Compose([ # transforms.Resize((256, 256)), # Resize according to your model's input size transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # Apply transformations and add batch dimension image = transform(image).unsqueeze(0) return image class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x) class Down(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x) class Up(nn.Module): def __init__(self, in_channels, out_channels, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) else: self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 = self.up(x1) # Resize x1 to match x2 diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x = torch.cat([x2, x1], dim=1) return self.conv(x) class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, n_channels=3, n_classes=1, bilinear=True): super().__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear # Encoder self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) factor = 2 if bilinear else 1 self.down4 = Down(512, 1024 // factor) # Decoder self.up1 = Up(1024, 512 // factor, bilinear) self.up2 = Up(512, 256 // factor, bilinear) self.up3 = Up(256, 128 // factor, bilinear) self.up4 = Up(128, 64, bilinear) self.outc = OutConv(64, n_classes) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) logits = self.outc(x) return torch.sigmoid(logits) def init_weights(self): # Initialize with Kaiming initialization def init_fn(m): if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') self.apply(init_fn) def load_UNet_model(model_path): print(f"Loading model from {model_path}") model = torch.load(model_path, weights_only=False, map_location=torch.device('cpu')) # Load the model (entire model saved with torch.save) model.eval() # Set the model to evaluation mode return model def predict(selected_model, uploaded_image): if selected_model == "Instance Segmentation Model (U-Net)": print("Predicting using U-Net") model_path = "UNet_Model.pth" # Path to your trained model else: print("Invalid model selected") return None # Visualize predictions (call visualize_predictions) return visualize_predictions(uploaded_image, model_path) # Visualization function for contours and IoU def visualize_predictions(uploaded_image, model_path="UNet.pth"): model = load_UNet_model(model_path) image = process_image(uploaded_image) # Make prediction with torch.no_grad(): output = model(image) binary_pred = (output > 0.5).float().cpu().numpy() # Prediction as a binary mask pred_prob = output.squeeze().cpu().numpy() # Prediction probabilities (for heatmap) # Visualization part (assumes ground truth is available) fig, axes = plt.subplots(1, 4, figsize=(16, 4)) # Original image img = np.array(uploaded_image) / 255.0 # Normalize the image to [0, 1] axes[0].imshow(img) axes[0].set_title('Original Image') axes[0].axis('off') # Ground truth (this is just an example, you should provide the actual mask) # For the sake of demonstration, we use a dummy mask ground_truth = np.zeros_like(binary_pred[0, 0]) axes[1].imshow(ground_truth, cmap='gray') axes[1].set_title('Ground Truth') axes[1].axis('off') # Prediction Probability axes[2].imshow(pred_prob, cmap='jet', vmin=0, vmax=1) axes[2].set_title('Prediction Probability') axes[2].axis('off') # Calculate IoU (Intersection over Union) intersection = np.logical_and(binary_pred[0, 0] > 0.5, ground_truth > 0.5).sum() union = np.logical_or(binary_pred[0, 0] > 0.5, ground_truth > 0.5).sum() iou = intersection / union if union > 0 else 0 axes[3].imshow(img) contours, _ = cv2.findContours(binary_pred[0, 0].astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contour_img = np.zeros_like(binary_pred[0, 0]) cv2.drawContours(contour_img, contours, -1, 1, 2) # Add the contour overlay with IoU text axes[3].imshow(contour_img, cmap='Reds', alpha=0.5) axes[3].set_title(f'Prediction Contour') axes[3].axis('off') plt.tight_layout() # Save the figure to a BytesIO object and return it as an image buf = io.BytesIO() plt.savefig(buf, format='png') buf.seek(0) img = Image.open(buf) return img # Connect buttons to navigation functions start_trial_button.click( fn=go_to_system_page, inputs=None, # Pass the current page state outputs=[welcome_page, system_page] ) back_button.click( fn=go_to_welcome_page, inputs=None, # Pass the current page state outputs=[welcome_page, system_page] ) upload_image_button.click( fn=predict, inputs=[radio, img_input], outputs=img_output ) demo.launch(share=True)