Update app.py

app.py

@@ -3,36 +3,40 @@ from urllib.request import urlopen
 from PIL import Image
 import timm
 import torch
+import time
 
-# Load the model
 model = timm.create_model("hf_hub:Marqo/nsfw-image-detection-384", pretrained=True)
 model = model.eval()
 
-# Prepare the data transformation
 data_config = timm.data.resolve_model_data_config(model)
 transforms = timm.data.create_transform(**data_config, is_training=False)
 
-# Prediction function
 def predict(image):
+    start_time = time.time()
     with torch.no_grad():
-        # Transform the image
         input_tensor = transforms(image).unsqueeze(0)
-        # Run the model
         output = model(input_tensor).softmax(dim=-1).cpu()
-        # Get class names
         class_names = model.pretrained_cfg["label_names"]
-        # Create the result dictionary
         result = {class_names[i]: float(output[0, i]) for i in range(len(class_names))}
-    return result
+    end_time = time.time()
+    inference_time = end_time - start_time
+    return result, f"Inference time: {inference_time:.2f} seconds"
 
-# Gradio interface
-interface = gr.Interface(
+
+demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=3),
+    inputs=gr.Image(type="pil", height=512),
+    outputs=[
+        gr.Label(num_top_classes=2),
+        gr.Textbox(label="Inference Time")
+    ],
     title="NSFW Image Detection",
-    description="Upload an image to detect if it is NSFW or Safe for Work."
+    description=(
+        "Upload an image to detect if it is **NSFW (Not Safe For Work)** or **Safe For Work (SFW)**.\n\n"
+        "This app uses the [Marqo/nsfw-image-detection-384](https://huggingface.co/Marqo/nsfw-image-detection-384) "
+        "image classification model from Hugging Face's `timm` library."
+    )
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()