Update app.py

app.py

@@ -76,7 +76,45 @@ advanced_css = f"""
     background: linear-gradient(90deg, {theme_color} 0%, #C06C84 100%);
 }}
 """
+# Define CNN model
+class Classifier(nn.Module):
+    def __init__(self):
+        super(Classifier, self).__init__()
+        self.cnn_layers = resnet18(weights=ResNet18_Weights.DEFAULT)
+        self.fc_layers = nn.Sequential(
+            nn.Linear(1000, 512),
+            nn.Dropout(0.3),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Linear(128, 5),
+        )
 
+    def forward(self, x):
+        x = self.cnn_layers(x)
+        x = self.fc_layers(x)
+        return x
+
+# Pre-process
+preprocess = transforms.Compose([
+    transforms.Resize(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Load model
+model = Classifier()
+model.load_state_dict(torch.load("classify_nsfw_v3.0.pth", map_location="cpu"))
+model.eval()
+
+def predict(image_path):
+    img = Image.open(image_path).convert("RGB")
+    img = preprocess(img).unsqueeze(0)
+
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(img)[0], dim=0)
+
+    result = {labels[i]: float(prediction[i]) for i in range(5)}
+    return result
 with gr.Blocks(theme=gr.themes.Soft(), css=advanced_css) as demo:
     # 标题区
     with gr.Column(elem_classes="header-section"):