Posture / app.py
Pengi5659's picture
Update app.py
d32de0a verified
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()