app.py

import cv2
import numpy as np
from PIL import Image
import torch
from torchvision import models, transforms
from ultralytics import YOLO
import gradio as gr
import torch.nn as nn
import os 

# --- DOCUMENTATION STRINGS (English and Urdu) ---

USAGE_GUIDELINES = """
## 1. Quick Start Guide: Run Instructions

**English:**
1. **Upload:** Click the 'Upload Rice Image' box and select your image (JPG or PNG).
2. **Run:** Click the **"Run Analysis"** button.
3. **Review:** The output image will show each detected rice grain marked with its predicted class label.

**Urdu (اردو):**
## 1. فوری استعمال کی ہدایات: تجزیہ شروع کریں
1. **تصویر اپ لوڈ کریں:** 'چاول کی تصویر اپ لوڈ کریں' کے باکس پر کلک کریں اور اپنی تصویر (JPG یا PNG) منتخب کریں۔
2. **چلائیں:** **"تجزیہ شروع کریں"** بٹن پر کلک کریں۔
3. **نتائج دیکھیں:** آؤٹ پٹ تصویر ہر شناخت شدہ چاول کے دانے کو اس کی پیش گوئی کردہ قسم کے لیبل کے ساتھ دکھائے گی۔
"""

INPUT_EXPLANATION = """
## 2. Expected Inputs / متوقع ان پٹ

| Input Field (ان پٹ فیلڈ) | Purpose (مقصد) | Requirement (ضرورت) |
| :--- | :--- | :--- |
| **Upload Image** / تصویر اپ لوڈ کریں | The image containing the rice grains for analysis. | Must be a single image file (JPG, PNG). The input should preferably contain clear, separated rice grains. |

**Important Note (اہم نوٹ):** For the best detection and classification accuracy, ensure the rice grains are scattered (not heavily overlapping) and the background is simple.
"""

OUTPUT_EXPLANATION = """
## 3. Expected Outputs (Detection and Classification) / متوقع آؤٹ پٹ

The output is the original image overlayed with results from the two-stage AI model:

*   **Bounding Boxes:** Each individual rice grain detected by the YOLO model is enclosed in a **Green rectangle**.
*   **Labels:** Above each rectangle, the predicted class label is displayed:
    *   **c9** (A specific rice variety)
    *   **kant** (Another specific rice variety)
    *   **superf** (A third specific rice variety)

**Urdu (اردو):**
آؤٹ پٹ اصل تصویر ہوگی جس پر دو مراحل پر مشتمل AI ماڈل کے نتائج لگائے گئے ہیں:
*   **نشان زد ڈبے:** YOLO ماڈل سے تلاش کیے گئے ہر چاول کے دانے کے گرد **سبز رنگ کا مستطیل** لگایا جائے گا۔
*   **لیبلز:** ہر مستطیل کے اوپر اس کی پیش گوئی کردہ قسم کا لیبل درج ہوگا: c9، kant، یا superf۔
"""

# --- CORE LOGIC (KEPT AS IS) ---

# Initialize device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load models
# NOTE: These paths (best.pt, rice_resnet_model.pth) must exist in the execution directory.
try:
    yolo_model = YOLO('best.pt') 
    resnet = models.resnet50(weights=None) # Use weights=None since we are loading custom weights
    resnet.fc = nn.Linear(resnet.fc.in_features, 3)
    resnet.load_state_dict(torch.load('rice_resnet_model.pth', map_location=device))
    resnet = resnet.to(device)
    resnet.eval()
except Exception as e:
    print(f"Error loading models: {e}. Ensure 'best.pt' and 'rice_resnet_model.pth' are present.")
    yolo_model = None
    resnet = None
    
# Class labels
class_labels = ["c9", "kant", "superf"]

# Image transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def classify_crop(crop_img):
    """Classify a single rice grain crop"""
    if resnet is None:
        return "Error"
    image = transform(crop_img).unsqueeze(0).to(device)
    with torch.no_grad():
        output = resnet(image)
        _, predicted = torch.max(output, 1)
    return class_labels[predicted.item()]

def detect_and_classify(input_image):
    """Process the image and classify each grain"""
    if yolo_model is None or resnet is None:
        raise gr.Error("Models failed to load. Cannot proceed with analysis.")
        
    if input_image is None:
        raise gr.Error("Please upload an image or select an example.")
        
    image = np.array(input_image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # YOLO Detection
    results = yolo_model(image, verbose=False)[0]
    boxes = results.boxes.xyxy.cpu().numpy()

    if len(boxes) == 0:
        gr.Warning("No rice grains detected in the image.")
        return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))


    for box in boxes:
        x1, y1, x2, y2 = map(int, box[:4])
        crop = image[y1:y2, x1:x2]
        
        # Check if crop is valid
        if crop.shape[0] > 0 and crop.shape[1] > 0:
            crop_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
            predicted_label = classify_crop(crop_pil)
    
            # Draw results
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) # Green box
            cv2.putText(image, 
                       predicted_label, 
                       (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 
                       0.9, 
                       (36, 255, 12), 
                       2) # Green text

    return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

# --- GRADIO UI (Single Tab) ---

with gr.Blocks(title="Rice Variety Classification") as demo:
    
    gr.Markdown(
        """
        # Rice Variety Classification System / چاول کی اقسام کی شناخت کا نظام
        Upload an image containing rice grains. The system will detect and classify each grain using a two-stage AI pipeline (YOLO for detection, ResNet for classification).
        """
    )
    
    # 1. GUIDELINES SECTION
    with gr.Accordion(" Tips & Guidelines (ہدایات)", open=False):
        gr.Markdown(USAGE_GUIDELINES)
        gr.Markdown("---")
        gr.Markdown(INPUT_EXPLANATION)
        gr.Markdown("---")
        gr.Markdown(OUTPUT_EXPLANATION)
        
    gr.Markdown("---")

    # 2. APPLICATION INTERFACE
    gr.Markdown("## Start Analysis / تجزیہ شروع کریں")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Step 1: Upload an Image of Rice /  چاول کی تصویر اپ لوڈ کریں")
            image_input = gr.Image(type="pil")
            gr.Markdown("## Step 2: Click Run Analysis /تجزیہ شروع کریں  ")
            submit_btn = gr.Button("Run Analysis / تجزیہ شروع کریں", variant="primary")
            gr.Markdown("## Outputs / نتائج ")
            output_image = gr.Image(scale=2, interactive=True)
    
    submit_btn.click(
        fn=detect_and_classify,
        inputs=image_input,
        outputs=output_image
    )
    
    gr.Markdown("---")

    # 3. EXAMPLES SECTION
    gr.Markdown("## Example Images / مثال تصاویر")
    # Ensure "samples/" directory exists and contains these images
    gr.Examples(
        examples=[
            "samples/rice1.jpg",
            "samples/rice2.jpg",
            "samples/rice3.jpg",
            "samples/rice4.jpg",
            "samples/rice5.jpg",
            "samples/rice6.jpg"
        ],
        inputs=image_input,
        outputs=output_image, # Required for proper caching and execution
        fn=detect_and_classify, # Required for proper caching and execution
        # cache_examples=True,
        label="Click to load and run a sample image / نمونہ تصویر لوڈ اور رن کرنے کے لیے کلک کریں"
    )

demo.queue()
demo.launch()