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
from datetime import datetime
import json

# ============================================
# RICE ANALYZER PRO - Advanced Analytics Dashboard
# A professional grain analysis platform with statistics and batch processing
# ============================================

# --- SYSTEM CONFIGURATION ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Model initialization
try:
    detection_model = YOLO('best.pt') 
    classifier_network = models.resnet50(weights=None)
    classifier_network.fc = nn.Linear(classifier_network.fc.in_features, 3)
    classifier_network.load_state_dict(torch.load('rice_resnet_model.pth', map_location=device))
    classifier_network = classifier_network.to(device)
    classifier_network.eval()
    models_loaded = True
except Exception as e:
    print(f" Model initialization failed: {e}")
    detection_model = None
    classifier_network = None
    models_loaded = False

# Variety definitions
VARIETY_MAP = {
    0: "C9 Premium",
    1: "Kant Special",
    2: "Superfine Grade"
}

VARIETY_COLORS = {
    "C9 Premium": (255, 100, 100),      # Red
    "Kant Special": (100, 100, 255),    # Blue
    "Superfine Grade": (100, 255, 100)  # Green
}

# Data preprocessing pipeline
image_preprocessor = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# --- ANALYTICS FUNCTIONS ---

def analyze_single_grain(grain_image):
    """Perform classification on an individual grain"""
    if not models_loaded:
        return "System Error"
    
    tensor_input = image_preprocessor(grain_image).unsqueeze(0).to(device)
    with torch.no_grad():
        prediction = classifier_network(tensor_input)
        class_idx = torch.argmax(prediction, dim=1).item()
    return VARIETY_MAP[class_idx]

def compute_distribution_stats(variety_counts):
    """Generate statistical summary of grain distribution"""
    total = sum(variety_counts.values())
    if total == 0:
        return "No grains detected for analysis."
    
    stats = [" **Distribution Analysis**\n"]
    stats.append(f" Total Grains Detected: **{total}**\n")
    stats.append("\n**Breakdown by Variety:**\n")
    
    for variety, count in sorted(variety_counts.items(), key=lambda x: x[1], reverse=True):
        percentage = (count / total) * 100
        bar_length = int(percentage / 5)
        bar = "█" * bar_length + "░" * (20 - bar_length)
        stats.append(f"- {variety}: {count} grains ({percentage:.1f}%) {bar}\n")
    
    # Dominant variety
    dominant = max(variety_counts.items(), key=lambda x: x[1])
    stats.append(f"\n **Dominant Variety:** {dominant[0]}")
    
    return "".join(stats)

def execute_batch_analysis(input_img):
    """Main processing pipeline with analytics"""
    if not models_loaded:
        raise gr.Error(" Analysis engine unavailable. Please check model files.")
    
    if input_img is None:
        raise gr.Error(" Please upload an image to begin analysis.")
    
    # Convert and prepare image
    img_array = np.array(input_img)
    img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    
    # Detection phase
    detection_results = detection_model(img_bgr, verbose=False)[0]
    bounding_boxes = detection_results.boxes.xyxy.cpu().numpy()
    
    if len(bounding_boxes) == 0:
        gr.Warning("⚠️ No rice grains found. Try a clearer image.")
        return (
            Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)),
            "No grains detected in the provided image."
        )
    
    # Classification phase with tracking
    variety_counter = {v: 0 for v in VARIETY_MAP.values()}
    
    for box in bounding_boxes:
        x1, y1, x2, y2 = map(int, box[:4])
        grain_crop = img_bgr[y1:y2, x1:x2]
        
        if grain_crop.shape[0] > 0 and grain_crop.shape[1] > 0:
            grain_pil = Image.fromarray(cv2.cvtColor(grain_crop, cv2.COLOR_BGR2RGB))
            variety_label = analyze_single_grain(grain_pil)
            variety_counter[variety_label] += 1
            
            # Visualization with color coding
            color = VARIETY_COLORS[variety_label]
            cv2.rectangle(img_bgr, (x1, y1), (x2, y2), color, 3)
            
            # Label with background
            label_text = variety_label
            (text_w, text_h), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
            cv2.rectangle(img_bgr, (x1, y1-text_h-10), (x1+text_w, y1), color, -1)
            cv2.putText(img_bgr, label_text, (x1, y1-5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    
    # Generate analytics
    statistics_report = compute_distribution_stats(variety_counter)
    
    return (
        Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)),
        statistics_report
    )

# --- GRADIO INTERFACE ---

custom_css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.header-box {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 30px;
    border-radius: 15px;
    color: white;
    text-align: center;
    margin-bottom: 20px;
}
.stat-box {
    background: #f8f9fa;
    border-left: 4px solid #667eea;
    padding: 15px;
    border-radius: 8px;
    margin: 10px 0;
}
"""

with gr.Blocks(css=custom_css, title="Rice Analyzer Pro") as app:
    
    gr.HTML("""
        <div class="header-box">
            <h1>🌾 Rice Analyzer Pro</h1>
            <p style="font-size: 18px; margin-top: 10px;">
                Advanced Grain Analytics Platform | Professional Quality Assessment
            </p>
        </div>
    """)
    
    with gr.Tabs():
        # Main Analysis Tab
        with gr.Tab("Analysis app"):
            gr.Markdown("""
            ### How It Works
            1. **Upload** your rice sample image (clear photos work best)
            2. **Analyze** Click on "Start Analysis" and let Our AI model detect and classify each grain
            3. **Review** detailed statistics and visual results
            
            **Color Coding:** 🔵 C9 Premium | 🔴 Kant Special | 🟢 Superfine Grade
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    image_upload = gr.Image(type="pil", label="📸 Sample Image")
                    analyze_button = gr.Button(" Start Analysis", variant="primary", size="lg")
                
                with gr.Column(scale=1):
                    output_visualization = gr.Image(label=" Annotated Results")
            
            with gr.Row():
                statistics_output = gr.Markdown(label=" Statistical Report")
            
            analyze_button.click(
                fn=execute_batch_analysis,
                inputs=image_upload,
                outputs=[output_visualization, statistics_output]
            )
        
        # Documentation Tab
        with gr.Tab(" Documentation"):
            gr.Markdown("""
            ## System Overview
            
            **Rice Analyzer Pro** uses a two-stage deep learning pipeline:
            - **Stage 1:** YOLO-based grain detection
            - **Stage 2:** ResNet50 variety classification
            
            ### Supported Varieties
            | Variety | Description | Market Grade |
            |---------|-------------|--------------|
            | C9 Premium | High-quality long grain | Premium |
            | Kant Special | Medium grain specialty | Standard |
            | Superfine Grade | Ultra-refined grain | Super Fine |
            
            ### Best Practices
            - Use well-lit images with minimal shadows
            - Ensure grains are separated (not clumped)
            - Plain backgrounds yield better results
            - Recommended resolution: 1024x1024 or higher
            
            ### Technical Specifications
            - Detection Model: YOLOv8
            - Classification: ResNet50 (Fine-tuned)
            - Processing: GPU-accelerated (when available)
            """)
    
    gr.Markdown("---")
    gr.Markdown("### Sample Gallery")
    
    gr.Examples(
        examples=[
            "samples/rice1.jpg",
            "samples/rice2.jpg",
            "samples/rice4.jpg",
            "samples/rice5.jpg"
        ],
        inputs=image_upload,
        outputs=[output_visualization, statistics_output],
        fn=execute_batch_analysis,
        label="Click any sample to run instant analysis"
    )

if __name__ == "__main__":
    app.queue()
    app.launch()