Spaces:
Sleeping
Sleeping
| 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() | |