Fix utils.py

app.py

@@ -1,11 +1,9 @@
 import gradio as gr
-import torch
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 import pandas as pd
-from PIL import Image, ImageDraw
+from PIL import Image
 import matplotlib.pyplot as plt
-import matplotlib
 
 from pipelines.detection.yolo_v8 import Yolov8Pipeline
 from pipelines.detection.yolo_stamp import YoloStampPipeline
@@ -41,17 +39,21 @@ def doc_predict(image, det_choice, seg_choice, emb_choice):
         cropped_stamp = image.crop(box.tolist())
         segmented_stamps.append(dlv3(cropped_stamp) if seg_choice else cropped_stamp)
 
-    widths, heights = zip(*(i.size for i in segmented_stamps))
+    if len(segmented_stamps) != 0:
+        widths, heights = zip(*(i.size for i in segmented_stamps))
 
-    total_width = sum(widths)
-    max_height = max(heights)
+        total_width = sum(widths)
+        max_height = max(heights)
 
-    concatenated_stamps = Image.new('RGB', (total_width, max_height))
+        concatenated_stamps = Image.new('RGB', (total_width, max_height))
 
-    x_offset = 0
-    for im in segmented_stamps:
-        concatenated_stamps.paste(im, (x_offset,0))
-        x_offset += im.size[0]
+        x_offset = 0
+        for im in segmented_stamps:
+            concatenated_stamps.paste(im, (x_offset,0))
+            x_offset += im.size[0]
+
+    else: 
+        concatenated_stamps = Image.new('RGB', (0, 0))
 
     embeddings = []
     if emb_choice == 'vits8':

utils.py

@@ -1,3 +1,7 @@
+import matplotlib
+import matplotlib.pyplot as plt
+from PIL import Image, ImageDraw
+
 def heatmap(data, row_labels, col_labels, ax=None,
             cbar_kw=None, cbarlabel="", **kwargs):
     """