Update app.py

app.py

@@ -8,6 +8,8 @@ import numpy as np
 import kagglehub
 from PIL import Image
 from glob import glob
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend
 import matplotlib.pyplot as plt
 from matplotlib import patches
 from torchvision import transforms as T
@@ -16,6 +18,7 @@ import shutil
 import tempfile
 from pathlib import Path
 import json
+from io import BytesIO
 
 # Try to import spaces for Hugging Face Spaces GPU support
 try:
@@ -90,7 +93,6 @@ class Visualization:
             self.im_paths[data_type] = im_paths
     
     def plot_single(self, im_path, bboxes):
-        fig, ax = plt.subplots(figsize=(8, 8))
         or_im = np.array(Image.open(im_path).convert("RGB"))
         height, width, _ = or_im.shape
 
@@ -102,16 +104,19 @@ class Visualization:
             x_max = int((x_center + w / 2) * width)
             y_max = int((y_center + h / 2) * height)
             
-            color = (random.randint(0, 255)/255, random.randint(0, 255)/255, random.randint(0, 255)/255)
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
             cv2.rectangle(img=or_im, pt1=(x_min, y_min), pt2=(x_max, y_max), 
-                         color=(int(color[0]*255), int(color[1]*255), int(color[2]*255)), thickness=3)
+                         color=color, thickness=3)
         
-        ax.imshow(or_im)
-        ax.axis("off")
-        ax.set_title(f"Number of objects: {len(bboxes)}")
-        plt.tight_layout()
+        # Add text overlay
+        cv2.putText(or_im, f"Objects: {len(bboxes)}", (10, 30), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+        
+        # Convert BGR to RGB if needed
+        if len(or_im.shape) == 3 and or_im.shape[2] == 3:
+            or_im = cv2.cvtColor(or_im, cv2.COLOR_BGR2RGB)
         
-        return fig
+        return Image.fromarray(or_im)
 
     def vis_samples(self, data_type, n_samples=4):
         if data_type not in self.vis_datas:
@@ -156,7 +161,14 @@ class Visualization:
                    ha='center', va='bottom', fontsize=10, color='navy')
         
         plt.tight_layout()
-        return fig
+        
+        # Convert matplotlib figure to PIL Image
+        fig.canvas.draw()
+        img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+        img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        plt.close(fig)
+        
+        return Image.fromarray(img_array)
 
 def download_dataset():
     """Download the dataset using kagglehub"""
@@ -429,7 +441,7 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
             data_type_analysis = gr.Dropdown(["train", "valid", "test"], value="train", label="Dataset Type")
             analyze_btn = gr.Button("Analyze Distribution")
         
-        distribution_plot = gr.Plot(label="Class Distribution")
+        distribution_plot = gr.Image(label="Class Distribution", type="pil")
         analysis_status = gr.Textbox(label="Status", interactive=False)
         
         analyze_btn.click(analyze_class_distribution, inputs=data_type_analysis,