import gradio as gr
import torch
import torch.nn as nn
import json
import os

class SoilFertilityModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(input_size, 32)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(32, num_classes)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# تحميل الموديل
config_path = "config.json"
with open(config_path) as f:
    config = json.load(f)

model = SoilFertilityModel(config["input_size"], config["num_classes"])
model.load_state_dict(torch.load("soil_fertility_model.pth", map_location="cpu"))
model.eval()

def predict(N, P, K, ph, ec, oc, S, zn, fe, cu, Mn, B):
    inputs = torch.tensor([[N, P, K, ph, ec, oc, S, zn, fe, cu, Mn, B]], dtype=torch.float32)
    with torch.no_grad():
        outputs = model(inputs)
        _, pred = torch.max(outputs, 1)
        return int(pred.item())

inputs = [gr.Number(label=name) for name in ["N", "P", "K", "ph", "ec", "oc", "S", "zn", "fe", "cu", "Mn", "B"]]
outputs = gr.Label(num_top_classes=3, label="Soil Fertility Class")

demo = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title="Soil Fertility Classifier")
demo.launch()