Update app.py

app.py

@@ -1,32 +1,22 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from PIL import Image
-import requests
-from io import BytesIO
-import numpy as np
+from transformers import pipeline
 
-# Load the pre-trained model and tokenizer
-model_name = "distilbert/distilbert-base-uncased"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# Load the DETR object detection pipeline
+object_detection_pipeline = pipeline("object-detection", model="facebook/detr-resnet-50")
 
-# Function to preprocess the image
-def preprocess_image(image):
-    image = Image.open(BytesIO(image))
-    image = image.resize((256, 256))  # Resize the image to match the model's input size
-    return np.array(image)
+# Function to perform object detection on an image
+def detect_objects(image):
+    # Perform object detection
+    results = object_detection_pipeline(image)
+    # Extract bounding boxes and object labels
+    bounding_boxes = [obj["bbox"] for obj in results]
+    labels = [obj["label"] for obj in results]
+    # Return bounding boxes and labels
+    return bounding_boxes, labels
 
-# Function to make predictions
-def classify_image(image):
-    image = preprocess_image(image)
-    inputs = tokenizer(image, return_tensors="pt", padding=True, truncation=True)
-    outputs = model(**inputs)
-    logits = outputs.logits.detach().numpy()[0]
-    probabilities = np.exp(logits) / np.exp(logits).sum(-1)
-    predicted_class = np.argmax(probabilities)
-    return {str(i): float(prob) for i, prob in enumerate(probabilities)}
+# Define Gradio interface
+inputs = gr.inputs.Image()
+outputs = gr.outputs.ObjectDetection(labels=["person", "car", "truck", "bicycle", "motorcycle"]) # Customize labels as needed
 
-# Create a Gradio interface
-input_image = gr.inputs.Image(shape=(256, 256))
-output_label = gr.outputs.Label(num_top_classes=3)
-gr.Interface(classify_image, inputs=input_image, outputs=output_label).launch()
+# Create Gradio interface
+gr.Interface(fn=detect_objects, inputs=inputs, outputs=outputs, title="Object Detection", description="Detect objects in images.").launch()