Update app.py

app.py

@@ -4,7 +4,6 @@ import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import resnet18, ResNet18_Weights
 from PIL import Image
-from io import BytesIO
 
 # number convert to label
 labels = ["Drawings", "Hentai", "Neutral", "Porn", "Sexy"]
@@ -59,17 +58,6 @@ def predict(image_path):
     result = {labels[i]: float(prediction[i]) for i in range(5)}
     return result
 
-# API 预测接口
-def api_predict(image: gr.Image) -> dict:
-    img = Image.open(BytesIO(image)).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
-
 # 美化 Gradio 界面
 css = """
     .gradio-container {
@@ -104,8 +92,7 @@ css = """
     footer {visibility: hidden}
 """
 
-# 使用 Blocks API 创建界面
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(theme="huggingface", css=css) as demo:
     gr.Markdown("# 🌟 NSFW 图片分类器")
     gr.Markdown(description)
 
@@ -123,10 +110,7 @@ with gr.Blocks(css=css) as demo:
             ]
             gr.Examples(examples=examples, inputs=image_input, label="点击上传示例图片")
 
-    submit_btn.click(predict, inputs=image_input, outputs=output_label)
-
-    # 手动为 API 设置路由
-    demo.add_api_route("/api/predict", api_predict, methods=["POST"])
+    submit_btn.click(predict, inputs=image_input, outputs=output_label, api_name="predict")
 
 # 启动 Web 界面
 demo.launch()