import gradio as gr
from ultralytics import YOLO
import os
import json
import tempfile
from pathlib import Path

# load the fixed model
print("Loading model...")
model = YOLO('/app/best_fixed.pt', task='classify')
print("Model loaded! Classes:", model.names)

CLASS_NAMES = {
    0: 'Armyworm', 1: 'Grasshopper', 2: 'aphids',
    3: 'bean_rust', 4: 'beans_angular_leaf_spot',
    5: 'beans_anthracnose', 6: 'beans_healthy',
    7: 'cutworm', 8: 'maize_common_rust',
    9: 'maize_healthy', 10: 'maize_leaf_blight',
    11: 'maize_streak_virus', 12: 'potato healthy',
    13: 'potato_early_blight', 14: 'potato_late_blight',
    15: 'rice_bacterial_leaf_blight', 16: 'rice_blast',
    17: 'rice_brown_spot', 18: 'rice_healthy',
    19: 'stem_borer', 20: 'thrips', 21: 'weevil',
}

def predict(image):
    try:
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
            image.save(tmp.name)
            tmp_path = tmp.name

        results = model.predict(tmp_path, task='classify', verbose=False)
        os.unlink(tmp_path)

        probs = results[0].probs
        top_class_idx = int(probs.top1)
        confidence = float(probs.top1conf)
        label = CLASS_NAMES.get(top_class_idx, 'Unknown')

        return json.dumps({
            'label': label,
            'score': round(confidence, 4),
            'success': True
        })

    except Exception as e:
        return json.dumps({
            'label': 'error',
            'score': 0.0,
            'success': False,
            'error': str(e)
        })

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type='pil', label='Upload crop image'),
    outputs=gr.Text(label='Prediction result'),
    title='AgriVision Crop Disease & Pest Classifier',
    description='Upload a crop image to detect diseases and pests.',
)

if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0', server_port=7860)