Update utils.py

utils.py

@@ -90,20 +90,15 @@ class UnionFind:
         return labels
 
 def visualise_single_image_prediction(image_as_np_array, predictions, filename):
-    figure, subplot = plt.subplots(1, 1, figsize=(10, 10))
+    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))
     subplot.imshow(image_as_np_array)
     plot_bboxes(subplot, predictions["panels"], color="green")
-    plot_bboxes(subplot, predictions["texts"], color="red", visibility=predictions["is_essential_text"])
+    plot_bboxes(subplot, predictions["texts"], color="red", add_index=True)
     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
@@ -146,24 +141,14 @@ 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")
-    
-    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.plot([x1, x2], [y1, y2], color="red", linewidth=2, linestyle="dashed", alpha=score)
 
     subplot.axis("off")
     if filename is not None:
@@ -174,18 +159,17 @@ def visualise_single_image_prediction(image_as_np_array, predictions, filename):
     plt.close()
     return image
 
-def plot_bboxes(subplot, bboxes, color="red", visibility=None):
-    if visibility is None:
-        visibility = [1] * len(bboxes)
+def plot_bboxes(subplot, bboxes, color="red", add_index=False):
     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)
@@ -334,44 +318,27 @@ def merge_overlapping_ranges(ranges):
     return merged_ranges
 
 def sort_text_boxes_in_reading_order(text_bboxes, sorted_panel_bboxes):
-    """
-    Sort text boxes in reading order (left-to-right, top-to-bottom) within panels.
-
-    Args:
-        text_bboxes: List of text bounding boxes [x1, y1, x2, y2].
-        sorted_panel_bboxes: List of panel bounding boxes, assumed to be sorted in reading order.
-
-    Returns:
-        List of indices of text boxes sorted in reading order.
-    """
     text_bboxes = convert_to_list_of_lists(text_bboxes)
     sorted_panel_bboxes = convert_to_list_of_lists(sorted_panel_bboxes)
 
     if len(text_bboxes) == 0:
         return []
 
-    # Step 1: Map each text box to its panel
-    panel_id_for_text = get_text_to_panel_mapping(text_bboxes, sorted_panel_bboxes)
-
-    # Step 2: Sort text boxes first by their panel and then within each panel
-    text_indices = list(range(len(text_bboxes)))
-    # Sort by panel ID first
-    text_indices.sort(key=lambda i: panel_id_for_text[i])
-
-    # Step 3: Group by panel ID and sort within each panel
-    from itertools import groupby
-
-    sorted_indices = []
-    for panel_id, group in groupby(text_indices, key=lambda i: panel_id_for_text[i]):
-        group_indices = list(group)
-        text_bboxes_within_panel = [text_bboxes[i] for i in group_indices]
-
-        # Sort within the panel (left-to-right, top-to-bottom)
-        sorted_within_panel = sort_texts_within_panel(text_bboxes_within_panel)
-        sorted_indices.extend([group_indices[i] for i in sorted_within_panel])
-
-    return sorted_indices
+    def indices_of_same_elements(nums):
+        groups = groupby(range(len(nums)), key=lambda i: nums[i])
+        return [list(indices) for _, indices in groups]
 
+    panel_id_for_text = get_text_to_panel_mapping(text_bboxes, sorted_panel_bboxes)
+    indices_of_texts = list(range(len(text_bboxes)))
+    indices_of_texts, panel_id_for_text = zip(*sorted(zip(indices_of_texts, panel_id_for_text), key=lambda x: x[1]))
+    indices_of_texts = list(indices_of_texts)
+    grouped_indices = indices_of_same_elements(panel_id_for_text)
+    for group in grouped_indices:
+        subset_of_text_indices = [indices_of_texts[i] for i in group]
+        text_bboxes_of_subset = [text_bboxes[i] for i in subset_of_text_indices]
+        sorted_subset_indices = sort_texts_within_panel(text_bboxes_of_subset)
+        indices_of_texts[group[0] : group[-1] + 1] = [subset_of_text_indices[i] for i in sorted_subset_indices]
+    return indices_of_texts
 
 def get_text_to_panel_mapping(text_bboxes, sorted_panel_bboxes):
     text_to_panel_mapping = []
@@ -394,10 +361,24 @@ def get_text_to_panel_mapping(text_bboxes, sorted_panel_bboxes):
     return text_to_panel_mapping
 
 def sort_texts_within_panel(rects):
-    # Sort primarily by Y-coordinate (top-to-bottom), then X-coordinate (left-to-right)
-    sorted_indices = sorted(range(len(rects)), key=lambda i: (rects[i][1], rects[i][0]))
-    return sorted_indices
+    smallest_y = float("inf")
+    smallest_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)
+    
+    reference_point = Point(smallest_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
+    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
 
 def x1y1wh_to_x1y1x2y2(bbox):
     x1, y1, w, h = bbox