MooRagab/Planet_Detection

1import gradio as gr2from PIL import Image3import torch4import torch.nn.functional as F5from torchvision import transforms6import matplotlib.pyplot as plt7import os8 9device="cuda" if torch.cuda.is_available() else "cpu"10 11model = torch.load("best_model.pth", map_location=torch.device("cpu"))12model.eval()13model.to(device)14 15 16class_labels = ['Apple___Apple_scab',17 'Apple___Black_rot',18 'Apple___Cedar_apple_rust',19 'Apple___healthy',20 'Blueberry___healthy',21 'Cherry_(including_sour)___Powdery_mildew',22 'Cherry_(including_sour)___healthy',23 'Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot',24 'Corn_(maize)___Common_rust_',25 'Corn_(maize)___Northern_Leaf_Blight',26 'Corn_(maize)___healthy',27 'Grape___Black_rot',28 'Grape___Esca_(Black_Measles)',29 'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)',30 'Grape___healthy',31 'Orange___Haunglongbing_(Citrus_greening)',32 'Peach___Bacterial_spot',33 'Peach___healthy',34 'Pepper,_bell___Bacterial_spot',35 'Pepper,_bell___healthy',36 'Potato___Early_blight',37 'Potato___Late_blight',38 'Potato___healthy',39 'Raspberry___healthy',40 'Soybean___healthy',41 'Squash___Powdery_mildew',42 'Strawberry___Leaf_scorch',43 'Strawberry___healthy',44 'Tomato___Bacterial_spot',45 'Tomato___Early_blight',46 'Tomato___Late_blight',47 'Tomato___Leaf_Mold',48 'Tomato___Septoria_leaf_spot',49 'Tomato___Spider_mites Two-spotted_spider_mite',50 'Tomato___Target_Spot',51 'Tomato___Tomato_Yellow_Leaf_Curl_Virus',52 'Tomato___Tomato_mosaic_virus',53 'Tomato___healthy']54 55 56transform_valid = transforms.Compose([57    transforms.Resize((256, 256)),58    transforms.CenterCrop(224),59    transforms.ToTensor(),60    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])61])62 63 64def predict_image(image):65    image = Image.fromarray(image)66    image_tensor = transform_valid(image).unsqueeze(0).to(device)67    68    model.eval()69    with torch.no_grad():70        output = model(image_tensor)71        probabilities = F.softmax(output, dim=1)72        print(probabilities)73        prob, pred = torch.max(probabilities, 1)74    75    class_name = class_labels[pred.item()]76    probability = prob.item() * 10077 78    79    fig, ax = plt.subplots(figsize=(5, 5))80    ax.imshow(image)81    ax.axis('off')82    ax.set_title(f"Predicted: {class_name}\nProbability: {probability:.2f}%")83 84    confidence_message = (85        "This image is likely not a plant leaf." if probability < 5086        else "Prediction made with confidence."87    )88 89    return fig, confidence_message90 91example_list = [["examples/" + example] for example in os.listdir("examples")]92 93 94interface = gr.Interface(95    fn=predict_image,96    inputs=gr.Image(type="numpy", label="Upload Plant Image"),97    outputs=[98        gr.Plot(label="Prediction Result"),99        gr.Textbox(label="Confidence Message")100    ],101    live=True,102    title="Plant Disease Detection",103    description="Upload an image of a plant leaf, and the model will predict the plant species and its disease with the predicted probability.",104    examples=example_list105)106 107interface.launch(share=True)