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

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

# Load models
yolo_model = YOLO('best.pt')  # Make sure this file is uploaded
resnet = models.resnet50(pretrained=False)
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()

# 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):
    """ایک چاول کے دانے کو درجہ بند کریں"""
    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):
    """تصویر پر کارروائی کریں اور ہر دانے کو شناخت کریں"""
    image = np.array(input_image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    results = yolo_model(image)[0]
    boxes = results.boxes.xyxy.cpu().numpy()

    for box in boxes:
        x1, y1, x2, y2 = map(int, box[:4])
        crop = image[y1:y2, x1:x2]
        crop_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
        predicted_label = classify_crop(crop_pil)

        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, 
                   predicted_label, 
                   (x1, y1-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 
                   0.9, 
                   (36, 255, 12), 
                   2)

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

with gr.Blocks(title="چاول کی اقسام کی درجہ بندی") as demo:
    gr.Markdown("""
    ## 🍚 چاول کی اقسام کی شناخت کا نظام  
    ایک تصویر اپ لوڈ کریں جس میں چاول کے دانے ہوں۔  
    سسٹم ہر دانے کو شناخت اور درجہ بند کرے گا۔
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="چاول کی تصویر اپ لوڈ کریں")
            submit_btn = gr.Button("تجزیہ شروع کریں", variant="primary")
        with gr.Column():
            output_image = gr.Image(label="نتائج", interactive=False)
    
    submit_btn.click(
        fn=detect_and_classify,
        inputs=image_input,
        outputs=output_image
    )
    
    # ✅ Move this block inside the `with gr.Blocks(...)` scope
    gr.Examples(
        examples=[
            "samples/rice1.jpg",
            "samples/rice2.jpg",
            "samples/rice3.jpg",
            "samples/rice4.jpg",
            "samples/rice5.jpg",
            "samples/rice6.jpg"
        ],
        inputs=image_input,
        label="مثال تصاویر"
    )
    gr.Markdown("""
### ℹ️ ہدایات:
- ✅ واضح اور الگ الگ چاول کے دانے والی تصویر اپ لوڈ کریں۔
- ⚠️ اگر دانے آپس میں جُڑے ہوں یا ایک دوسرے پر چڑھے ہوں، تو نتائج متاثر ہو سکتے ہیں۔
- 📸 بہتر پہچان کے لیے تصویر کا پس منظر صاف اور دانے منتشر (پھیلے ہوئے) ہونے چاہئیں۔
- 🖼️ آپ اوپر دی گئی مثال تصاویر کو بھی دیکھ سکتے ہیں۔
""")



demo.launch()