import torch
import gradio as gr
from transformers import ConvNextForImageClassification, AutoImageProcessor
from PIL import Image

# Load model configuration and weights manually
model = ConvNextForImageClassification.from_pretrained("facebook/convnext-base-224")  # Load the base model
model.load_state_dict(torch.load("convnext_base_finetuned.pth", map_location="cpu"))  # Load your finetuned weights
model.eval()

# Load the processor
processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224")

# Define a function to predict the class from an image
def predict(image):
    # Preprocess the image
    inputs = processor(images=image, return_tensors="pt")
    
    # Perform inference
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_class = torch.argmax(outputs.logits, dim=1).item()
    
    return predicted_class

# Create Gradio interface for user input
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Textbox())

iface.launch()