Update app.py

app.py

@@ -1,24 +1,28 @@
-# prompt: gradio image 分类 not safe for work
+import torch
 from PIL import Image
+from transformers import AutoModelForImageClassification, ViTImageProcessor
 import gradio as gr
-from transformers import pipeline
-# Load the image classification pipeline
-classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+
+# Load the model and processor
+model_name = "Falconsai/nsfw_image_detection"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+processor = ViTImageProcessor.from_pretrained(model_name)
 
 # Define a function to classify the image and return the results
 def classify_image(img):
-    # Convert the Gradio image input to a PIL image
     pil_image = Image.fromarray(img.astype('uint8'), 'RGB')
-    # Classify the image using the pipeline
-    results = classifier(pil_image)
-    # Format the results for display in Gradio
-    formatted_results = {result['label']: result['score'] for result in results}
-    return formatted_results
+    inputs = processor(images=pil_image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_label = logits.argmax(-1).item()
+    label = model.config.id2label[predicted_label]
+    return label
 
 # Create the Gradio interface
 image_input = gr.inputs.Image(shape=(256, 256))
-label_output = gr.outputs.Label(num_top_classes=3)
+label_output = gr.outputs.Label()
 interface = gr.Interface(fn=classify_image, inputs=image_input, outputs=label_output)
 
 # Launch the interface
-interface.launch()
+interface.launch()