Update utils.py

utils.py

@@ -90,15 +90,20 @@ class UnionFind:
         return labels
 
 def visualise_single_image_prediction(image_as_np_array, predictions, filename):
-    h, w = image_as_np_array.shape[:2]
-    if h > w:
-        figure, subplot = plt.subplots(1, 1, figsize=(10, 10 * h / w))
-    else:
-        figure, subplot = plt.subplots(1, 1, figsize=(10 * w / h, 10))
+    figure, subplot = plt.subplots(1, 1, figsize=(10, 10))
     subplot.imshow(image_as_np_array)
     plot_bboxes(subplot, predictions["panels"], color="green")
-    plot_bboxes(subplot, predictions["texts"], color="red", add_index=True)
+    plot_bboxes(subplot, predictions["texts"], color="red", visibility=predictions["is_essential_text"])
     plot_bboxes(subplot, predictions["characters"], color="blue")
+    plot_bboxes(subplot, predictions["tails"], color="purple")
+
+    for i, name in enumerate(predictions["character_names"]):
+        char_bbox = predictions["characters"][i]
+        x1, y1, x2, y2 = char_bbox
+        subplot.text(x1, y1 - 2, name, 
+            verticalalignment='bottom', horizontalalignment='left',
+            bbox=dict(facecolor='blue', alpha=1, edgecolor='none'),  # Background settings
+            color='white', fontsize=8)
 
     COLOURS = [
         "#b7ff51", # green
@@ -141,14 +146,24 @@ def visualise_single_image_prediction(image_as_np_array, predictions, filename):
             subplot.plot([x2], [y2], color=random_colour, marker="o", markersize=5)
     
     for (i, j) in predictions["text_character_associations"]:
-        score = predictions["dialog_confidences"][i]
         bbox_i = predictions["texts"][i]
         bbox_j = predictions["characters"][j]
+        if not predictions["is_essential_text"][i]:
+            continue
         x1 = bbox_i[0] + (bbox_i[2] - bbox_i[0]) / 2
         y1 = bbox_i[1] + (bbox_i[3] - bbox_i[1]) / 2
         x2 = bbox_j[0] + (bbox_j[2] - bbox_j[0]) / 2
         y2 = bbox_j[1] + (bbox_j[3] - bbox_j[1]) / 2
-        subplot.plot([x1, x2], [y1, y2], color="red", linewidth=2, linestyle="dashed", alpha=score)
+        subplot.plot([x1, x2], [y1, y2], color="red", linewidth=2, linestyle="dashed")
+    
+    for (i, j) in predictions["text_tail_associations"]:
+        bbox_i = predictions["texts"][i]
+        bbox_j = predictions["tails"][j]
+        x1 = bbox_i[0] + (bbox_i[2] - bbox_i[0]) / 2
+        y1 = bbox_i[1] + (bbox_i[3] - bbox_i[1]) / 2
+        x2 = bbox_j[0] + (bbox_j[2] - bbox_j[0]) / 2
+        y2 = bbox_j[1] + (bbox_j[3] - bbox_j[1]) / 2
+        subplot.plot([x1, x2], [y1, y2], color="purple", linewidth=2, linestyle="dashed")
 
     subplot.axis("off")
     if filename is not None:
@@ -159,17 +174,18 @@ def visualise_single_image_prediction(image_as_np_array, predictions, filename):
     plt.close()
     return image
 
-def plot_bboxes(subplot, bboxes, color="red", add_index=False):
+def plot_bboxes(subplot, bboxes, color="red", visibility=None):
+    if visibility is None:
+        visibility = [1] * len(bboxes)
     for id, bbox in enumerate(bboxes):
+        if visibility[id] == 0:
+            continue
         w = bbox[2] - bbox[0]
         h = bbox[3] - bbox[1]
         rect = patches.Rectangle(
             bbox[:2], w, h, linewidth=1, edgecolor=color, facecolor="none", linestyle="solid"
         )
         subplot.add_patch(rect)
-        if add_index:
-            cx, cy = bbox[0] + w / 2, bbox[1] + h / 2
-            subplot.text(cx, cy, str(id), color=color, fontsize=10, ha="center", va="center")
 
 def sort_panels(rects):
     before_rects = convert_to_list_of_lists(rects)
@@ -362,20 +378,19 @@ def get_text_to_panel_mapping(text_bboxes, sorted_panel_bboxes):
 
 def sort_texts_within_panel(rects):
     smallest_y = float("inf")
-    smallest_x = float("inf")
+    greatest_x = float("-inf")
     for i, rect in enumerate(rects):
         x1, y1, x2, y2 = rect
         smallest_y = min(smallest_y, y1)
-        smallest_x = min(smallest_x, x1)
+        greatest_x = max(greatest_x, x2)
     
-    reference_point = Point(smallest_x, smallest_y)
+    reference_point = Point(greatest_x, smallest_y)
 
     polygons_and_index = []
     for i, rect in enumerate(rects):
         x1, y1, x2, y2 = rect
         polygons_and_index.append((box(x1,y1,x2,y2), i))
-    
-    # sort points by distance to this new reference point
+    # sort points by closest to reference point
     polygons_and_index = sorted(polygons_and_index, key=lambda x: reference_point.distance(x[0]))
     indices = [x[1] for x in polygons_and_index]
     return indices