Spaces:
Sleeping
Sleeping
File size: 3,310 Bytes
1622277 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
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()
|