Upload 3 files

app.py

@@ -0,0 +1,57 @@
+import gradio as gr
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import torch.nn as nn
+import os
+
+# ✅ Define Lightweight CNN Model (Same as trained)
+class SmallCNN(nn.Module):
+    def __init__(self):
+        super(SmallCNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)  # Reduced filters
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)  # Reduced filters
+        self.fc1 = nn.Linear(32 * 8 * 8, 10)  # 10-class classification (CIFAR-10)
+
+    def forward(self, x):
+        x = torch.relu(self.conv1(x))
+        x = torch.max_pool2d(x, 2)
+        x = torch.relu(self.conv2(x))
+        x = torch.max_pool2d(x, 2)
+        x = x.view(x.size(0), -1)
+        x = self.fc1(x)
+        return x
+
+# ✅ Load the trained model from Hugging Face Space
+model_path = os.path.join(os.getenv("SPACE_ROOT", ""), "light_cnn_model.pth")
+
+# ✅ Initialize model and load weights
+model = SmallCNN()
+model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+model.eval()
+
+# ✅ Define Image Transformation
+transform = transforms.Compose([
+    transforms.Resize((32, 32)),  # Resize image to 32x32 pixels
+    transforms.ToTensor(),  # Convert image to tensor
+])
+
+# ✅ Define Prediction Function
+def predict(image):
+    image = transform(image).unsqueeze(0)  # Convert image to tensor and add batch dimension
+    with torch.no_grad():
+        output = model(image)  # Forward pass through model
+        prediction = torch.argmax(output, dim=1).item()  # Get predicted class
+    return f"Predicted Class: {prediction}"
+
+# ✅ Create Gradio Interface
+interface = gr.Interface(
+    fn=predict,  # Function to process image
+    inputs=gr.Image(type="pil"),  # User uploads an image
+    outputs="text",  # Model returns a text output
+    title="Lightweight CNN Image Classification",
+    description="Upload an image to classify using the trained CNN model.",
+)
+
+# ✅ Launch the Gradio App
+interface.launch()

light_cnn_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2afca200cc840345889a7174a85ab796df3ebf4ae681522f8ed43df4b39ed756
+size 104984

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+torch
+torchvision
+Pillow