Update app.py

app.py

@@ -11,51 +11,55 @@ processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-101')
 model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
 
 def object_detection(image, confidence_threshold):
-    # Convert the input to a PIL Image object if it's not already
-    if not isinstance(image, Image.Image):
-        image = Image.open(io.BytesIO(image))
-    
-    # Preprocess the image
-    inputs = processor(images=image, return_tensors="pt")
+    try:
+        # Convert the input to a PIL Image object if it's not already
+        if not isinstance(image, Image.Image):
+            image = Image.open(io.BytesIO(image))
+        
+        # Preprocess the image
+        inputs = processor(images=image, return_tensors="pt")
 
-    # Perform object detection
-    outputs = model(**inputs)
+        # Perform object detection
+        outputs = model(**inputs)
 
-    # Extract bounding boxes and labels
-    target_sizes = torch.tensor([image.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=confidence_threshold)[0]
+        # Extract bounding boxes and labels
+        target_sizes = torch.tensor([image.size[::-1]])
+        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=confidence_threshold)[0]
 
-    # Plot the image with bounding boxes
-    plt.figure(figsize=(16, 10))
-    plt.imshow(image)
-    ax = plt.gca()
+        # Plot the image with bounding boxes
+        plt.figure(figsize=(16, 10))
+        plt.imshow(image)
+        ax = plt.gca()
 
-    detected_objects = []
-    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        box = [round(i, 2) for i in box.tolist()]
-        xmin, ymin, xmax, ymax = box
-        width, height = xmax - xmin, ymax - ymin
+        detected_objects = []
+        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+            box = [round(i, 2) for i in box.tolist()]
+            xmin, ymin, xmax, ymax = box
+            width, height = xmax - xmin, ymax - ymin
 
-        ax.add_patch(plt.Rectangle((xmin, ymin), width, height, fill=False, color='red', linewidth=3))
-        text = f'{model.config.id2label[label.item()]}: {round(score.item(), 3)}'
-        ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
-        detected_objects.append(text)
+            ax.add_patch(plt.Rectangle((xmin, ymin), width, height, fill=False, color='red', linewidth=3))
+            text = f'{model.config.id2label[label.item()]}: {round(score.item(), 3)}'
+            ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
+            detected_objects.append(text)
 
-    plt.axis('off')
+        plt.axis('off')
 
-    # Save the plot to an image buffer
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    plt.close()
+        # Save the plot to an image buffer
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png')
+        buf.seek(0)
+        plt.close()
 
-    # Convert buffer to an Image object
-    result_image = Image.open(buf)
+        # Convert buffer to an Image object
+        result_image = Image.open(buf)
 
-    # Join detected objects into a single string
-    detected_objects_text = "\n".join(detected_objects)
+        # Join detected objects into a single string
+        detected_objects_text = "\n".join(detected_objects)
 
-    return result_image, detected_objects_text
+        return result_image, detected_objects_text
+    
+    except Exception as e:
+        return Image.new("RGB", (224, 224), color="white"), str(e)
 
 # Define the Gradio interface
 demo = gr.Interface(