Update app.py

app.py

@@ -1,7 +1,7 @@
 import cv2
-import torch
 import numpy as np
 from PIL import Image
+import torch
 from torchvision import models, transforms
 from ultralytics import YOLO
 import gradio as gr
@@ -11,10 +11,8 @@ import torch.nn as nn
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load models
-yolo_model = YOLO('best.pt')  # Make sure this file is uploaded to your Space
+yolo_model = YOLO('best.pt')  # Make sure this file is uploaded
 resnet = models.resnet50(pretrained=False)
-
-# Modify ResNet for 3 classes
 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)
@@ -38,12 +36,17 @@ def classify_crop(crop_img):
         _, predicted = torch.max(output, 1)
     return class_labels[predicted.item()]
 
-def detect_and_classify(image):
-    """Process full image with YOLO + ResNet"""
-    image = np.array(image)
+def detect_and_classify(input_image):
+    """Process uploaded image"""
+    # Convert Gradio Image to OpenCV format
+    image = np.array(input_image)
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    
+    # YOLO Detection
     results = yolo_model(image)[0]
     boxes = results.boxes.xyxy.cpu().numpy()
-
+    
+    # Process each detection
     for box in boxes:
         x1, y1, x2, y2 = map(int, box[:4])
         crop = image[y1:y2, x1:x2]
@@ -52,35 +55,36 @@ def detect_and_classify(image):
 
         # Draw bounding box and label
         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)
+        cv2.putText(image, 
+                   predicted_label, 
+                   (x1, y1-10), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 
+                   0.9, 
+                   (36, 255, 12), 
+                   2)
 
+    # Convert back to RGB for Gradio
     return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
 
-# Gradio Interface
-with gr.Blocks(title="چاول کا شناختی نظام") as demo:
+# Create Gradio interface
+with gr.Blocks(title="Rice Classification") as demo:
     gr.Markdown("""
-    # چاول کا شناختی نظام
-    ایک تصویر اپ لوڈ کریں جس میں چاول کے دانے ہوں۔ نظام ہر دانے کو پہچان کر اس کی قسم بتائے گا۔
+    ## 🍚 Rice Variety Classifier
+    Upload an image containing rice grains. The system will detect and classify each grain.
     """)
     
     with gr.Row():
-        input_image = gr.Image(type="pil", label="تصویر داخل کریں")
-        output_image = gr.Image(type="pil", label="نتیجہ")
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Rice Image")
+            submit_btn = gr.Button("Analyze", variant="primary")
+        with gr.Column():
+            output_image = gr.Image(label="Detection Results", interactive=False)
     
-    submit_btn = gr.Button("تشخیص کریں")
     submit_btn.click(
         fn=detect_and_classify,
-        inputs=input_image,
+        inputs=image_input,
         outputs=output_image
     )
-    
-    gr.Examples(
-        examples=[["example1.jpg"], ["example2.jpg"]],  # Add your example images
-        inputs=input_image,
-        outputs=output_image,
-        fn=detect_and_classify,
-        cache_examples=True
-    )
 
+# Launch the app
 demo.launch()