import gradio as gr
import numpy as np
import cv2
import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt

# Load Pretrained AI Model (You can replace this with your custom model)
model = torch.hub.load("pytorch/vision:v0.10.0", "resnet18", pretrained=True)
model.eval()

# Function to analyze deforestation using AI model
def analyze_deforestation(image):
    image = Image.fromarray(image).convert("RGB")
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    img_tensor = transform(image).unsqueeze(0)
    
    # AI Prediction
    with torch.no_grad():
        output = model(img_tensor)
    
    # Simulated deforestation risk (random value for demo)
    deforestation_risk = np.random.uniform(0.3, 0.9)  
    
    # Identify locations for seed-dropping (For simplicity, we use random dots)
    h, w = image.size
    num_seeds = int(deforestation_risk * 20)  # More risk = More seeds
    seed_locations = np.random.randint(0, min(h, w), (num_seeds, 2))

    # Plot Results
    plt.figure(figsize=(6, 6))
    plt.imshow(image)
    for loc in seed_locations:
        plt.scatter(loc[0], loc[1], color="red", s=20, marker="o")  # Mark seed locations
    plt.title(f"Deforestation Risk: {deforestation_risk:.2f} | Seeds: {num_seeds}")
    plt.axis("off")
    plt.savefig("output.png")
    
    return "output.png", f"Risk Level: {deforestation_risk:.2f} | Seeds Dropped: {num_seeds}"

# Create Gradio UI
interface = gr.Interface(
    fn=analyze_deforestation,
    inputs=gr.Image(type="numpy"),
    outputs=["image", "text"],
    title="🌱 Plant Me AI - Drone Seed Dispersal 🌍",
    description="Upload a satellite image of a forest area. The AI will analyze deforestation risks and predict seed-dropping locations for drone-based planting."
)

# Launch the web app
interface.launch(debug=True)