Update app.py

app.py

@@ -1,27 +1,23 @@
-# prompt: gradio image 分类
-import fastai
-from fastai.vision import *
-from fastai.vision.all import load_learner,PILImage
+# prompt: gradio image 分类 not safe for work
+from PIL import Image
 import gradio as gr
-from transformers import AutoModelForSequenceClassification
-# Load the model
+from transformers import pipeline
+# Load the image classification pipeline
+classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
 
-model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
-model = load_learner(model)
-
-# Define an image classification function
-def classify_image(image):
-    img = PILImage.create(image)
-
-    # Make a prediction
-    pred_class, pred_idx, probs = model.predict(img)
-    
-    # Return the prediction as a dictionary
-    return {model.dls.vocab[i]: float(probs[i]) for i in range(len(probs))}
+# 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
 
 # Create the Gradio interface
-image_input = gr.Image()
-label_output = gr.Label(num_top_classes=2)
+image_input = gr.inputs.Image(shape=(256, 256))
+label_output = gr.outputs.Label(num_top_classes=3)
 interface = gr.Interface(fn=classify_image, inputs=image_input, outputs=label_output)
 
 # Launch the interface