Update app.py

app.py

@@ -8,6 +8,7 @@ import torch
 import numpy as np
 import spaces
 import subprocess
+from io import BytesIO
 
 # Ensure flash-attn is installed correctly
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -78,13 +79,14 @@ def plot_image_with_bboxes(image_np, bboxes, labels=None):
     plt.axis('off')
     
     # Save the plot to a BytesIO object
-    from io import BytesIO
     buf = BytesIO()
     plt.savefig(buf, format='png')
     plt.close()
     buf.seek(0)
     
-    return buf
+    # Convert the BytesIO buffer to PIL Image
+    pil_image = Image.open(buf)
+    return pil_image
 
 # Gradio function to process uploaded images
 @spaces.GPU