Update app.py

app.py

@@ -31,6 +31,20 @@ def detect_objects(url: str):
     target_sizes = torch.tensor([image.size[::-1]])
     results = yolos_image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
 
+    # Draw bounding boxes on the image
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        image_draw = ImageDraw.Draw(image)
+        image_draw.rectangle(box.tolist(), outline="red", width=2)
+        image_draw.text((box[0], box[1]), f"{yolos_model.config.id2label[label.item()]}: {round(score.item(), 3)}", fill="red")
+
+    # Save the modified image to a byte stream
+    image_byte_array = io.BytesIO()
+    image.save(image_byte_array, format="PNG")
+
+    # Encode the byte stream as a base64 string
+    image_base64 = base64.b64encode(image_byte_array.getvalue()).decode("utf-8")
+
+    
     # Format and return the results
     detected_objects = []
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
@@ -41,4 +55,7 @@ def detect_objects(url: str):
             "location": box
         })
 
-    return {"detected_objects": detected_objects}
+    return {"detected_objects": detected_objects, "image_base64": image_base64}
+
+except Exception as e:
+    raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")