import torch
import numpy as np
import random
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from shapely.geometry import Point, box
import networkx as nx
from copy import deepcopy
from itertools import groupby
def move_to_device(inputs, device):
if hasattr(inputs, "keys"):
return {k: move_to_device(v, device) for k, v in inputs.items()}
elif isinstance(inputs, list):
return [move_to_device(v, device) for v in inputs]
elif isinstance(inputs, tuple):
return tuple([move_to_device(v, device) for v in inputs])
elif isinstance(inputs, np.ndarray):
return torch.from_numpy(inputs).to(device)
else:
return inputs.to(device)
class UnionFind:
def __init__(self, n):
self.parent = list(range(n))
self.size = [1] * n
self.num_components = n
@classmethod
def from_adj_matrix(cls, adj_matrix):
ufds = cls(adj_matrix.shape[0])
for i in range(adj_matrix.shape[0]):
for j in range(adj_matrix.shape[1]):
if adj_matrix[i, j] > 0:
ufds.unite(i, j)
return ufds
@classmethod
def from_adj_list(cls, adj_list):
ufds = cls(len(adj_list))
for i in range(len(adj_list)):
for j in adj_list[i]:
ufds.unite(i, j)
return ufds
@classmethod
def from_edge_list(cls, edge_list, num_nodes):
ufds = cls(num_nodes)
for edge in edge_list:
ufds.unite(edge[0], edge[1])
return ufds
def find(self, x):
if self.parent[x] == x:
return x
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
def unite(self, x, y):
x = self.find(x)
y = self.find(y)
if x != y:
if self.size[x] < self.size[y]:
x, y = y, x
self.parent[y] = x
self.size[x] += self.size[y]
self.num_components -= 1
def get_components_of(self, x):
x = self.find(x)
return [i for i in range(len(self.parent)) if self.find(i) == x]
def are_connected(self, x, y):
return self.find(x) == self.find(y)
def get_size(self, x):
return self.size[self.find(x)]
def get_num_components(self):
return self.num_components
def get_labels_for_connected_components(self):
map_parent_to_label = {}
labels = []
for i in range(len(self.parent)):
parent = self.find(i)
if parent not in map_parent_to_label:
map_parent_to_label[parent] = len(map_parent_to_label)
labels.append(map_parent_to_label[parent])
return labels
def visualise_single_image_prediction(image_as_np_array, predictions, filename):
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", 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
"#f50a8f", # pink
"#4b13b6", # purple
"#ddaa34", # orange
"#bea2a2", # brown
]
colour_index = 0
character_cluster_labels = predictions["character_cluster_labels"]
unique_label_sorted_by_frequency = sorted(list(set(character_cluster_labels)), key=lambda x: character_cluster_labels.count(x), reverse=True)
for label in unique_label_sorted_by_frequency:
root = None
others = []
for i in range(len(predictions["characters"])):
if character_cluster_labels[i] == label:
if root is None:
root = i
else:
others.append(i)
if colour_index >= len(COLOURS):
random_colour = COLOURS[0]
while random_colour in COLOURS:
random_colour = "#" + "".join([random.choice("0123456789ABCDEF") for j in range(6)])
else:
random_colour = COLOURS[colour_index]
colour_index += 1
bbox_i = predictions["characters"][root]
x1 = bbox_i[0] + (bbox_i[2] - bbox_i[0]) / 2
y1 = bbox_i[1] + (bbox_i[3] - bbox_i[1]) / 2
subplot.plot([x1], [y1], color=random_colour, marker="o", markersize=5)
for j in others:
# draw line from centre of bbox i to centre of bbox j
bbox_j = predictions["characters"][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=random_colour, linewidth=2)
subplot.plot([x2], [y2], color=random_colour, marker="o", markersize=5)
for (i, j) in predictions["text_character_associations"]:
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.axis("off")
if filename is not None:
plt.savefig(filename, bbox_inches="tight", pad_inches=0)
figure.canvas.draw()
image = np.array(figure.canvas.renderer._renderer)
plt.close()
return image
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)
def sort_panels(rects):
before_rects = convert_to_list_of_lists(rects)
# slightly erode all rectangles initially to account for imperfect detections
rects = [erode_rectangle(rect, 0.05) for rect in before_rects]
G = nx.DiGraph()
G.add_nodes_from(range(len(rects)))
for i in range(len(rects)):
for j in range(len(rects)):
if i == j:
continue
if is_there_a_directed_edge(i, j, rects):
G.add_edge(i, j, weight=get_distance(rects[i], rects[j]))
else:
G.add_edge(j, i, weight=get_distance(rects[i], rects[j]))
while True:
cycles = sorted(nx.simple_cycles(G))
cycles = [cycle for cycle in cycles if len(cycle) > 1]
if len(cycles) == 0:
break
cycle = cycles[0]
edges = [e for e in zip(cycle, cycle[1:] + cycle[:1])]
max_cyclic_edge = max(edges, key=lambda x: G.edges[x]["weight"])
G.remove_edge(*max_cyclic_edge)
return list(nx.topological_sort(G))
def is_strictly_above(rectA, rectB):
x1A, y1A, x2A, y2A = rectA
x1B, y1B, x2B, y2B = rectB
return y2A < y1B
def is_strictly_below(rectA, rectB):
x1A, y1A, x2A, y2A = rectA
x1B, y1B, x2B, y2B = rectB
return y2B < y1A
def is_strictly_left_of(rectA, rectB):
x1A, y1A, x2A, y2A = rectA
x1B, y1B, x2B, y2B = rectB
return x2A < x1B
def is_strictly_right_of(rectA, rectB):
x1A, y1A, x2A, y2A = rectA
x1B, y1B, x2B, y2B = rectB
return x2B < x1A
def intersects(rectA, rectB):
return box(*rectA).intersects(box(*rectB))
def is_there_a_directed_edge(a, b, rects):
rectA = rects[a]
rectB = rects[b]
centre_of_A = [rectA[0] + (rectA[2] - rectA[0]) / 2, rectA[1] + (rectA[3] - rectA[1]) / 2]
centre_of_B = [rectB[0] + (rectB[2] - rectB[0]) / 2, rectB[1] + (rectB[3] - rectB[1]) / 2]
if np.allclose(np.array(centre_of_A), np.array(centre_of_B)):
return box(*rectA).area > (box(*rectB)).area
copy_A = [rectA[0], rectA[1], rectA[2], rectA[3]]
copy_B = [rectB[0], rectB[1], rectB[2], rectB[3]]
while True:
if is_strictly_above(copy_A, copy_B) and not is_strictly_left_of(copy_A, copy_B):
return 1
if is_strictly_above(copy_B, copy_A) and not is_strictly_left_of(copy_B, copy_A):
return 0
if is_strictly_right_of(copy_A, copy_B) and not is_strictly_below(copy_A, copy_B):
return 1
if is_strictly_right_of(copy_B, copy_A) and not is_strictly_below(copy_B, copy_A):
return 0
if is_strictly_below(copy_A, copy_B) and is_strictly_right_of(copy_A, copy_B):
return use_cuts_to_determine_edge_from_a_to_b(a, b, rects)
if is_strictly_below(copy_B, copy_A) and is_strictly_right_of(copy_B, copy_A):
return use_cuts_to_determine_edge_from_a_to_b(a, b, rects)
# otherwise they intersect
copy_A = erode_rectangle(copy_A, 0.05)
copy_B = erode_rectangle(copy_B, 0.05)
def get_distance(rectA, rectB):
return box(rectA[0], rectA[1], rectA[2], rectA[3]).distance(box(rectB[0], rectB[1], rectB[2], rectB[3]))
def use_cuts_to_determine_edge_from_a_to_b(a, b, rects):
rects = deepcopy(rects)
while True:
xmin, ymin, xmax, ymax = min(rects[a][0], rects[b][0]), min(rects[a][1], rects[b][1]), max(rects[a][2], rects[b][2]), max(rects[a][3], rects[b][3])
rect_index = [i for i in range(len(rects)) if intersects(rects[i], [xmin, ymin, xmax, ymax])]
rects_copy = [rect for rect in rects if intersects(rect, [xmin, ymin, xmax, ymax])]
# try to split the panels using a "horizontal" lines
overlapping_y_ranges = merge_overlapping_ranges([(y1, y2) for x1, y1, x2, y2 in rects_copy])
panel_index_to_split = {}
for split_index, (y1, y2) in enumerate(overlapping_y_ranges):
for i, index in enumerate(rect_index):
if y1 <= rects_copy[i][1] <= rects_copy[i][3] <= y2:
panel_index_to_split[index] = split_index
if panel_index_to_split[a] != panel_index_to_split[b]:
return panel_index_to_split[a] < panel_index_to_split[b]
# try to split the panels using a "vertical" lines
overlapping_x_ranges = merge_overlapping_ranges([(x1, x2) for x1, y1, x2, y2 in rects_copy])
panel_index_to_split = {}
for split_index, (x1, x2) in enumerate(overlapping_x_ranges[::-1]):
for i, index in enumerate(rect_index):
if x1 <= rects_copy[i][0] <= rects_copy[i][2] <= x2:
panel_index_to_split[index] = split_index
if panel_index_to_split[a] != panel_index_to_split[b]:
return panel_index_to_split[a] < panel_index_to_split[b]
# otherwise, erode the rectangles and try again
rects = [erode_rectangle(rect, 0.05) for rect in rects]
def erode_rectangle(bbox, erosion_factor):
x1, y1, x2, y2 = bbox
w, h = x2 - x1, y2 - y1
cx, cy = x1 + w / 2, y1 + h / 2
if w < h:
aspect_ratio = w / h
erosion_factor_width = erosion_factor * aspect_ratio
erosion_factor_height = erosion_factor
else:
aspect_ratio = h / w
erosion_factor_width = erosion_factor
erosion_factor_height = erosion_factor * aspect_ratio
w = w - w * erosion_factor_width
h = h - h * erosion_factor_height
x1, y1, x2, y2 = cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2
return [x1, y1, x2, y2]
def merge_overlapping_ranges(ranges):
"""
ranges: list of tuples (x1, x2)
"""
if len(ranges) == 0:
return []
ranges = sorted(ranges, key=lambda x: x[0])
merged_ranges = []
for i, r in enumerate(ranges):
if i == 0:
prev_x1, prev_x2 = r
continue
x1, x2 = r
if x1 > prev_x2:
merged_ranges.append((prev_x1, prev_x2))
prev_x1, prev_x2 = x1, x2
else:
prev_x2 = max(prev_x2, x2)
merged_ranges.append((prev_x1, prev_x2))
return merged_ranges
def sort_text_boxes_in_reading_order(text_bboxes, sorted_panel_bboxes):
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 []
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 = []
for text_bbox in text_bboxes:
shapely_text_polygon = box(*text_bbox)
all_intersections = []
all_distances = []
if len(sorted_panel_bboxes) == 0:
text_to_panel_mapping.append(-1)
continue
for j, annotation in enumerate(sorted_panel_bboxes):
shapely_annotation_polygon = box(*annotation)
if shapely_text_polygon.intersects(shapely_annotation_polygon):
all_intersections.append((shapely_text_polygon.intersection(shapely_annotation_polygon).area, j))
all_distances.append((shapely_text_polygon.distance(shapely_annotation_polygon), j))
if len(all_intersections) == 0:
text_to_panel_mapping.append(min(all_distances, key=lambda x: x[0])[1])
else:
text_to_panel_mapping.append(max(all_intersections, key=lambda x: x[0])[1])
return text_to_panel_mapping
def sort_texts_within_panel(rects):
smallest_y = float("inf")
greatest_x = float("-inf")
for i, rect in enumerate(rects):
x1, y1, x2, y2 = rect
smallest_y = min(smallest_y, y1)
greatest_x = max(greatest_x, x2)
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 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
def x1y1wh_to_x1y1x2y2(bbox):
x1, y1, w, h = bbox
return [x1, y1, x1 + w, y1 + h]
def x1y1x2y2_to_xywh(bbox):
x1, y1, x2, y2 = bbox
return [x1, y1, x2 - x1, y2 - y1]
def convert_to_list_of_lists(rects):
if isinstance(rects, torch.Tensor):
return rects.tolist()
if isinstance(rects, np.ndarray):
return rects.tolist()
return [[a, b, c, d] for a, b, c, d in rects]