Update app.py

app.py

@@ -4,6 +4,7 @@ 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"]
@@ -58,6 +59,17 @@ 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 {
@@ -92,7 +104,7 @@ css = """
     footer {visibility: hidden}
 """
 
-with gr.Blocks(theme="glass", css=css) as demo:
+with gr.Blocks(theme="huggingface", css=css) as demo:
     gr.Markdown("# 🌟 NSFW 图片分类器")
     gr.Markdown(description)
 
@@ -112,5 +124,8 @@ with gr.Blocks(theme="glass", css=css) as demo:
 
     submit_btn.click(predict, inputs=image_input, outputs=output_label)
 
+    # 启动 API 接口
+    gr.Interface(fn=api_predict, inputs=gr.Image(), outputs=gr.JSON(), api=True)
+
 # 启动 Web 界面
 demo.launch()