Update app.py

app.py

@@ -3,28 +3,19 @@ import torch
 import torchvision.transforms as transforms
 import torchvision.models as models
 from PIL import Image
-import os
-from torch.utils.data import DataLoader, Dataset
-from torchvision.datasets import ImageFolder
 
-# Define transformations for images
+# Load the ResNet18 model with pre-trained weights
+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()  # Set to evaluation mode
+
+# Define image transformation
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor()
 ])
 
-# Load dataset from your image folders
-dataset = ImageFolder(root="posture_samples", transform=transform)  # Ensure "posture_samples" contains two subfolders: "Good Posture-samples" and "Bad Posture-samples"
-
-# Create a DataLoader
-dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
-
-# Load pre-trained ResNet18 model
-model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
-model.fc = torch.nn.Linear(model.fc.in_features, 2)  # Adjust output for two classes
-model.eval()  # Set to evaluation mode
-
-# Define function to classify an image
+# Function to classify posture images
 def classify_image(image):
     image = transform(image).unsqueeze(0)
     output = model(image)
@@ -32,6 +23,7 @@ def classify_image(image):
     return "Good Posture" if predicted.item() == 0 else "Bad Posture"
 
 # Set up Gradio interface
-iface = gr.Interface(fn=classify_image, inputs=gr.Image(source="webcam", type="pil"), outputs="text")
+iface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil", tool="camera"), outputs="text")
 iface.launch()
 
+