import gradio as gr
import torch
import torchvision.transforms as transforms
import torchvision.models as models
from PIL import Image

# Load ResNet18 model
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.fc = torch.nn.Linear(model.fc.in_features, 2)  # Adjust for two classes
model.eval()

# Define image transformation
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

# Function to classify posture images
def classify_image(image):
    if image is None:
        return "No image provided!"
    
    image = transform(image).unsqueeze(0)
    output = model(image)
    _, predicted = torch.max(output, 1)
    return "Good Posture" if predicted.item() == 0 else "Bad Posture"

# Set up Gradio interface with live webcam capture
iface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil"), outputs="text")
iface.launch()