Update app.py

app.py

@@ -1,23 +1,25 @@
-import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+import gradio as gr
+import torch
+import torchvision.transforms as transforms
+import torchvision.models as models
+from PIL import Image
 
-import tensorflow as tf
-import cv2
-import imghdr
-import os
+# Load model (ensure it's uploaded to the Space)
+model = models.resnet18(pretrained=True)
+model.fc = torch.nn.Linear(model.fc.in_features, 2)  # Adjust for your classes
 
-# Clean image files
-data_dir = 'data'
-image_exts = ['jpeg', 'jpg', 'bmp', 'png']
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
+])
 
-for image_class in os.listdir(data_dir):
-    for image in os.listdir(os.path.join(data_dir, image_class)):
-        image_path = os.path.join(data_dir, image_class, image)
-        try:
-            img = cv2.imread(image_path)
-            tip = imghdr.what(image_path)
-            if tip not in image_exts:
-                print('Removing:', image_path)
-                os.remove(image_path)
-        except Exception as e:
-            print('Issue with image:', image_path)
+# Define the function to classify images
+def classify_image(image):
+    image = transform(image).unsqueeze(0)
+    output = model(image)
+    _, predicted = torch.max(output, 1)
+    return "Class_A" if predicted.item() == 0 else "Class_B"
+
+# Set up Gradio interface
+iface = gr.Interface(fn=classify_image, inputs="webcam", outputs="text")
+iface.launch()