Upload 2 files

app.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import gradio as gr
+from huggingface_hub import snapshot_download
+import os
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# === Model Definition ===
+class CNNModel(nn.Module):
+    def __init__(self, dropout_rate=0.5, hidden_size=512, use_batchnorm=True):
+        super(CNNModel, self).__init__()
+        self.use_batchnorm = use_batchnorm
+        
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(32) if use_batchnorm else nn.Identity()
+        
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(64) if use_batchnorm else nn.Identity()
+        
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm2d(128) if use_batchnorm else nn.Identity()
+        
+        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
+        self.bn4 = nn.BatchNorm2d(256) if use_batchnorm else nn.Identity()
+        
+        self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.bn5 = nn.BatchNorm2d(256) if use_batchnorm else nn.Identity()
+
+        self.pool = nn.MaxPool2d(2, 2)
+        self.dropout = nn.Dropout(dropout_rate)
+        
+        self.fc1 = nn.Linear(256 * 7 * 7, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, 1)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.bn1(self.conv1(x))))
+        x = self.pool(F.relu(self.bn2(self.conv2(x))))
+        x = self.pool(F.relu(self.bn3(self.conv3(x))))
+        x = self.pool(F.relu(self.bn4(self.conv4(x))))
+        x = self.pool(F.relu(self.bn5(self.conv5(x))))
+        x = torch.flatten(x, 1)
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
+
+# === Load Model from Hugging Face ===
+def load_model(repo_id="ZacToh/RandomSearchCNN"):
+    download_dir = snapshot_download(repo_id)
+    model_path = os.path.join(download_dir, "cnn_final_model.pth")
+
+    model = CNNModel()
+    checkpoint = torch.load(model_path, map_location=DEVICE)
+    state_dict = checkpoint.get("model_state_dict", checkpoint)
+    model.load_state_dict(state_dict, strict=False)
+
+    model.to(DEVICE)
+    model.eval()
+    return model
+
+model = load_model()
+
+# === Image Transform ===
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+])
+
+# === Prediction Function ===
+def predict(img: Image.Image) -> str:
+    img_tensor = transform(img).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        output = model(img_tensor)
+        prob = torch.sigmoid(output).item()
+    
+    label = "hf" if prob > 0.5 else "cc"
+    confidence = prob * 100 if label == "hf" else (1 - prob) * 100
+
+    return f"Prediction: {label.upper()} ({confidence:.2f}%)"
+
+# === Gradio UI ===
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="CNN Classifier: CC vs HF",
+    description="Upload an image to classify whether it belongs to the 'cc' or 'hf' category using a CNN model."
+).launch()

requirements.txt

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