Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from transformers import AutoImageProcessor, SiglipForImageClassification, pipeline | |
| from torchvision import transforms | |
| from PIL import Image | |
| import numpy as np | |
| import os | |
| # ------------------------------- | |
| # Model paths (local folders) | |
| # ------------------------------- | |
| hf_model_names = { | |
| "Rice": "models/Rice-Leaf-Disease", | |
| "Sugarcane": "models/sugarcane-plant-diseases-classification", | |
| "Tomato": "models/tomato-leaf-disease-classification-resnet50", | |
| "Corn/Wheat": "models/crop_leaf_diseases_vit" | |
| } | |
| # ------------------------------- | |
| # Utility: Load model offline or online | |
| # ------------------------------- | |
| def load_model_or_fallback(model_name, model_path, use_pipeline=False, skip_processor=False): | |
| if os.path.exists(model_path): | |
| print(f"✅ Loading local model: {model_path}") | |
| if use_pipeline: | |
| return pipeline("image-classification", model=model_path) | |
| elif skip_processor: | |
| model = SiglipForImageClassification.from_pretrained(model_path) | |
| return None, model | |
| else: | |
| processor = AutoImageProcessor.from_pretrained(model_path) | |
| model = SiglipForImageClassification.from_pretrained(model_path) | |
| return processor, model | |
| else: | |
| print(f"🌐 Model not found locally. Fetching from Hugging Face Hub: {model_name}") | |
| if use_pipeline: | |
| return pipeline("image-classification", model=model_name) | |
| elif skip_processor: | |
| model = SiglipForImageClassification.from_pretrained(model_name) | |
| return None, model | |
| else: | |
| processor = AutoImageProcessor.from_pretrained(model_name) | |
| model = SiglipForImageClassification.from_pretrained(model_name) | |
| return processor, model | |
| # ------------------------------- | |
| # Load models | |
| # ------------------------------- | |
| hf_processors = {} | |
| hf_models = {} | |
| # Rice | |
| hf_processors['Rice'], hf_models['Rice'] = load_model_or_fallback( | |
| "prithivMLmods/Rice-Leaf-Disease", hf_model_names["Rice"] | |
| ) | |
| # Sugarcane (skip processor) | |
| _, hf_models['Sugarcane'] = load_model_or_fallback( | |
| "dwililiya/sugarcane-plant-diseases-classification", | |
| hf_model_names["Sugarcane"], | |
| skip_processor=True | |
| ) | |
| # Tomato (pipeline) | |
| hf_models['Tomato'] = load_model_or_fallback( | |
| "wellCh4n/tomato-leaf-disease-classification-resnet50", | |
| hf_model_names["Tomato"], use_pipeline=True | |
| ) | |
| # Corn/Wheat (pipeline) | |
| hf_models['Corn/Wheat'] = load_model_or_fallback( | |
| "wambugu71/crop_leaf_diseases_vit", | |
| hf_model_names["Corn/Wheat"], use_pipeline=True | |
| ) | |
| # ------------------------------- | |
| # Sugarcane manual preprocessing | |
| # ------------------------------- | |
| sugarcane_transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| ]) | |
| # ------------------------------- | |
| # Disease mapping | |
| # ------------------------------- | |
| disease_dict = { | |
| "Rice": ["Bacterial Blight", "Blast", "Brown Spot", "Healthy", "Tungro"], | |
| "Sugarcane": ["Bacterial Blight", "Healthy", "Mosaic", "Red Rot", "Rust", "Yellow"], | |
| "Tomato": ["Early Blight", "Late Blight", "Healthy"], | |
| "Corn/Wheat": ["Healthy", "Rust", "Blight", "Leaf Spot"] # Adjust based on your model labels | |
| } | |
| # Remedies mapping | |
| remedies = { | |
| "Early Blight": "Remove infected leaves, apply fungicide.", | |
| "Late Blight": "Use fungicides and remove infected plants.", | |
| "Bacterial Blight": "Use resistant varieties and avoid overhead watering.", | |
| "Blast": "Use balanced fertilizer, apply fungicide.", | |
| "Brown Spot": "Ensure proper field drainage and avoid overcrowding.", | |
| "Tungro": "Control green leafhoppers and remove infected plants.", | |
| "Mosaic": "Remove infected plants, avoid spread.", | |
| "Red Rot": "Remove infected plants, apply fungicide.", | |
| "Rust": "Use fungicide and resistant varieties.", | |
| "Yellow": "Monitor plant, apply preventive measures.", | |
| "Leaf Spot": "Remove affected leaves and apply fungicide.", | |
| "Blight": "Use disease-free seeds and apply fungicides.", | |
| "Healthy": "No action needed." | |
| } | |
| # ------------------------------- | |
| # Prediction function | |
| # ------------------------------- | |
| def predict_disease(crop, img): | |
| if img is None: | |
| return "No image uploaded", "Please upload a leaf image." | |
| img_pil = Image.fromarray(img).convert("RGB") | |
| if crop == "Rice": | |
| inputs = hf_processors[crop](images=img_pil, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = hf_models[crop](**inputs) | |
| logits = outputs.logits | |
| probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist() | |
| predicted_idx = int(np.argmax(probs)) | |
| disease = disease_dict[crop][predicted_idx] | |
| advice = remedies.get(disease, "No advice available.") | |
| return disease, advice | |
| elif crop == "Sugarcane": | |
| img_tensor = sugarcane_transform(img_pil).unsqueeze(0) | |
| with torch.no_grad(): | |
| outputs = hf_models[crop](img_tensor) | |
| logits = outputs.logits | |
| probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist() | |
| predicted_idx = int(np.argmax(probs)) | |
| disease = disease_dict[crop][predicted_idx] | |
| advice = remedies.get(disease, "No advice available.") | |
| return disease, advice | |
| elif crop in ["Tomato", "Corn/Wheat"]: | |
| result = hf_models[crop](img_pil)[0] | |
| disease = result['label'] | |
| advice = remedies.get(disease, "No advice available.") | |
| return disease, advice | |
| else: | |
| return "Error", f"Model for {crop} is not available." | |
| # ------------------------------- | |
| # Gradio Interface | |
| # ------------------------------- | |
| custom_css = """ | |
| body, .gradio-container { | |
| background-image: url('https://media.istockphoto.com/id/1328004520/photo/healthy-young-soybean-crop-in-field-at-dawn.jpg?s=612x612&w=0&k=20&c=XRw20PArfhkh6LLgFrgvycPLm0Uy9y7lu9U7fLqabVY='); | |
| background-size: cover; | |
| background-repeat: no-repeat; | |
| background-attachment: fixed; | |
| min-height: 100vh !important; | |
| } | |
| .gradio-container > * { | |
| background-color: rgba(255, 255, 255, 0.88) !important; | |
| border-radius: 15px; | |
| padding: 20px; | |
| } | |
| """ | |
| with gr.Blocks(css=custom_css) as app: | |
| gr.Markdown("## 🌿 Crop Disease Detector") | |
| gr.Markdown("Upload a leaf image of your crop and get AI-based disease prediction with remedies.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| crop_input = gr.Dropdown(list(hf_model_names.keys()), label="Select Crop") | |
| img_input = gr.Image(type="numpy", label="Upload Leaf Image") | |
| predict_btn = gr.Button("🔍 Predict Disease") | |
| with gr.Column(): | |
| disease_output = gr.Textbox(label="Predicted Disease") | |
| advice_output = gr.Textbox(label="Recommended Action") | |
| predict_btn.click(predict_disease, inputs=[crop_input, img_input], outputs=[disease_output, advice_output]) | |
| # Launch | |
| app.launch(server_name="127.0.0.1", server_port=7860, share=True) | |