Update app.py

app.py

@@ -1,98 +1,98 @@
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torchvision import datasets, transforms
-from torch.utils.data import DataLoader
-from PIL import Image
-import gradio as gr
-import os
-
-# Device (CPU for compatibility with Hugging Face Spaces)
-device = torch.device("cpu")
-
-# Transform for training and uploaded images
-transform = transforms.Compose([
-    transforms.Resize((6, 6)),
-    transforms.ToTensor()
-])
-
-# Define a convolution block
-def conv(ic, oc):
-    ks=3
-    return nn.Sequential(
-        nn.Conv2d(ic, oc, stride=2, kernel_size=ks, padding=ks//2),
-        nn.BatchNorm2d(oc)
-    )
-
-# CNN Model
-class SimpleCNN(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.model = nn.Sequential(
-            conv(1, 8),
-            nn.Dropout2d(0.25),
-            nn.ReLU(),
-            conv(8, 16),
-            nn.Dropout2d(0.25),
-            nn.ReLU(),
-            conv(16, 10),
-            nn.Flatten()
-        )
-
-    def forward(self, x):
-        return self.model(x)
-
-# Training function
-def train_model():
-    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
-    batch_size = 36
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-
-    model = SimpleCNN().to(device)
-    optimizer = optim.Adam(model.parameters(), lr=0.005)
-    criterion = nn.CrossEntropyLoss()
-
-    model.train()
-    for epoch in range(3):  # Keep it light for HF Spaces
-        for images, labels in train_loader:
-            images, labels = images.to(device), labels.to(device)
-            optimizer.zero_grad()
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
-    return model
-
-# Load or train model
-model_path = "mnist_cnn.pt"
-if os.path.exists(model_path):
-    model = SimpleCNN().to(device)
-    model.load_state_dict(torch.load(model_path, map_location=device))
-else:
-    model = train_model()
-    torch.save(model.state_dict(), model_path)
-
-# Prediction function
-def predict(img):
-    if isinstance(img, Image.Image):
-        img = img.convert("L")
-    else:
-        return "Invalid image"
-    x = transform(img).unsqueeze(0).to(device)  # Shape: [1,1,8,8]
-    model.eval()
-    with torch.no_grad():
-        output = model(x)
-        pred = torch.argmax(output, dim=1).item()
-    return f"Predicted digit: {pred}"
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(shape=(28, 28), image_mode="L", invert_colors=True, sources=["upload", "canvas"]),
-    outputs="text",
-    title="MNIST Digit Classifier (6x6 CNN)",
-    description="Upload or draw a digit to classify it using a lightweight CNN trained on MNIST resized to 8×8."
-)
-
-if __name__ == "__main__":
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from PIL import Image
+import gradio as gr
+import os
+
+# Device (CPU for compatibility with Hugging Face Spaces)
+device = torch.device("cpu")
+
+# Transform for training and uploaded images
+transform = transforms.Compose([
+    transforms.Resize((6, 6)),
+    transforms.ToTensor()
+])
+
+# Define a convolution block
+def conv(ic, oc):
+    ks=3
+    return nn.Sequential(
+        nn.Conv2d(ic, oc, stride=2, kernel_size=ks, padding=ks//2),
+        nn.BatchNorm2d(oc)
+    )
+
+# CNN Model
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Sequential(
+            conv(1, 8),
+            nn.Dropout2d(0.25),
+            nn.ReLU(),
+            conv(8, 16),
+            nn.Dropout2d(0.25),
+            nn.ReLU(),
+            conv(16, 10),
+            nn.Flatten()
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+# Training function
+def train_model():
+    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+    batch_size = 36
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+
+    model = SimpleCNN().to(device)
+    optimizer = optim.Adam(model.parameters(), lr=0.005)
+    criterion = nn.CrossEntropyLoss()
+
+    model.train()
+    for epoch in range(3):  # Keep it light for HF Spaces
+        for images, labels in train_loader:
+            images, labels = images.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+    return model
+
+# Load or train model
+model_path = "mnist_cnn.pt"
+if os.path.exists(model_path):
+    model = SimpleCNN().to(device)
+    model.load_state_dict(torch.load(model_path, map_location=device))
+else:
+    model = train_model()
+    torch.save(model.state_dict(), model_path)
+
+# Prediction function
+def predict(img):
+    if isinstance(img, Image.Image):
+        img = img.convert("L")
+    else:
+        return "Invalid image"
+    x = transform(img).unsqueeze(0).to(device)  # Shape: [1,1,8,8]
+    model.eval()
+    with torch.no_grad():
+        output = model(x)
+        pred = torch.argmax(output, dim=1).item()
+    return f"Predicted digit: {pred}"
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(image_mode="L", invert_colors=True, sources=["upload", "canvas"]),
+    outputs="text",
+    title="MNIST Digit Classifier (6x6 CNN)",
+    description="Upload or draw a digit to classify it using a lightweight CNN trained on MNIST resized to 8×8."
+)
+
+if __name__ == "__main__":
     demo.launch()