from transformers import ConditionalDetrImageProcessor, TrOCRProcessor, ViTImageProcessor from transformers.image_transforms import center_to_corners_format import torch from typing import List from shapely.geometry import box from .utils import UnionFind, sort_panels, sort_text_boxes_in_reading_order, x1y1x2y2_to_xywh import numpy as np class MagiProcessor(): def __init__(self, config): self.config = config self.detection_image_preprocessor = None self.ocr_preprocessor = None self.crop_embedding_image_preprocessor = None if not config.disable_detections: assert config.detection_image_preprocessing_config is not None self.detection_image_preprocessor = ConditionalDetrImageProcessor.from_dict(config.detection_image_preprocessing_config) if not config.disable_ocr: assert config.ocr_pretrained_processor_path is not None self.ocr_preprocessor = TrOCRProcessor.from_pretrained(config.ocr_pretrained_processor_path) if not config.disable_crop_embeddings: assert config.crop_embedding_image_preprocessing_config is not None self.crop_embedding_image_preprocessor = ViTImageProcessor.from_dict(config.crop_embedding_image_preprocessing_config) def preprocess_inputs_for_detection(self, images, annotations=None): images = list(images) assert isinstance(images[0], np.ndarray) annotations = self._convert_annotations_to_coco_format(annotations) inputs = self.detection_image_preprocessor(images, annotations=annotations, return_tensors="pt") return inputs def preprocess_inputs_for_ocr(self, images): images = list(images) assert isinstance(images[0], np.ndarray) return self.ocr_preprocessor(images, return_tensors="pt").pixel_values def preprocess_inputs_for_crop_embeddings(self, images): images = list(images) assert isinstance(images[0], np.ndarray) return self.crop_embedding_image_preprocessor(images, return_tensors="pt").pixel_values def postprocess_detections_and_associations( self, predicted_bboxes, predicted_class_scores, original_image_sizes, get_character_character_matching_scores, get_text_character_matching_scores, get_dialog_confidence_scores, character_detection_threshold=0.3, panel_detection_threshold=0.2, text_detection_threshold=0.25, character_character_matching_threshold=0.7, text_character_matching_threshold=0.4, ): assert self.config.disable_detections is False batch_scores, batch_labels = predicted_class_scores.max(-1) batch_scores = batch_scores.sigmoid() batch_labels = batch_labels.long() batch_bboxes = center_to_corners_format(predicted_bboxes) # scale the bboxes back to the original image size if isinstance(original_image_sizes, List): img_h = torch.Tensor([i[0] for i in original_image_sizes]) img_w = torch.Tensor([i[1] for i in original_image_sizes]) else: img_h, img_w = original_image_sizes.unbind(1) scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(batch_bboxes.device) batch_bboxes = batch_bboxes * scale_fct[:, None, :] batch_panel_indices = self._get_indices_of_panels_to_keep(batch_scores, batch_labels, batch_bboxes, panel_detection_threshold) batch_character_indices = self._get_indices_of_characters_to_keep(batch_scores, batch_labels, batch_bboxes, character_detection_threshold) batch_text_indices = self._get_indices_of_texts_to_keep(batch_scores, batch_labels, batch_bboxes, text_detection_threshold) batch_character_character_matching_scores = get_character_character_matching_scores(batch_character_indices, batch_bboxes) batch_text_character_matching_scores = get_text_character_matching_scores(batch_text_indices, batch_character_indices) batch_dialog_confidence_scores = get_dialog_confidence_scores(batch_text_indices) # sort panels and texts in the reading order for batch_index in range(len(batch_scores)): panel_bboxes = batch_bboxes[batch_index][batch_panel_indices[batch_index]] panel_scores = batch_scores[batch_index][batch_panel_indices[batch_index]] text_bboxes = batch_bboxes[batch_index][batch_text_indices[batch_index]] text_scores = batch_scores[batch_index][batch_text_indices[batch_index]] sorted_panel_indices = sort_panels(panel_bboxes) batch_bboxes[batch_index][batch_panel_indices[batch_index]] = panel_bboxes[sorted_panel_indices] batch_scores[batch_index][batch_panel_indices[batch_index]] = panel_scores[sorted_panel_indices] sorted_panels = batch_bboxes[batch_index][batch_panel_indices[batch_index]] sorted_text_indices = sort_text_boxes_in_reading_order(text_bboxes, sorted_panels) batch_bboxes[batch_index][batch_text_indices[batch_index]] = text_bboxes[sorted_text_indices] batch_scores[batch_index][batch_text_indices[batch_index]] = text_scores[sorted_text_indices] batch_text_character_matching_scores[batch_index] = batch_text_character_matching_scores[batch_index][sorted_text_indices] batch_dialog_confidence_scores[batch_index] = batch_dialog_confidence_scores[batch_index][sorted_text_indices] results = [] for batch_index in range(len(batch_scores)): panel_bboxes = batch_bboxes[batch_index][batch_panel_indices[batch_index]] panel_scores = batch_scores[batch_index][batch_panel_indices[batch_index]] text_bboxes = batch_bboxes[batch_index][batch_text_indices[batch_index]] text_scores = batch_scores[batch_index][batch_text_indices[batch_index]] character_bboxes = batch_bboxes[batch_index][batch_character_indices[batch_index]] character_scores = batch_scores[batch_index][batch_character_indices[batch_index]] char_i, char_j = torch.where(batch_character_character_matching_scores[batch_index] > character_character_matching_threshold) character_character_associations = torch.stack([char_i, char_j], dim=1) text_boxes_to_match = batch_dialog_confidence_scores[batch_index] > text_character_matching_threshold if 0 in batch_text_character_matching_scores[batch_index].shape: text_character_associations = torch.zeros((0, 2), dtype=torch.long) else: most_likely_speaker_for_each_text = torch.argmax(batch_text_character_matching_scores[batch_index], dim=1)[text_boxes_to_match] text_indices = torch.arange(len(text_bboxes)).type_as(most_likely_speaker_for_each_text)[text_boxes_to_match] text_character_associations = torch.stack([text_indices, most_likely_speaker_for_each_text], dim=1) character_ufds = UnionFind.from_adj_matrix( batch_character_character_matching_scores[batch_index] > character_character_matching_threshold ) results.append({ "panels": panel_bboxes.tolist(), "panel_scores": panel_scores.tolist(), "texts": text_bboxes.tolist(), "text_scores": text_scores.tolist(), "characters": character_bboxes.tolist(), "character_scores": character_scores.tolist(), "character_character_associations": character_character_associations.tolist(), "text_character_associations": text_character_associations.tolist(), "character_cluster_labels": character_ufds.get_labels_for_connected_components(), "dialog_confidences": batch_dialog_confidence_scores[batch_index].tolist(), }) return results def postprocess_ocr_tokens(self, generated_ids, skip_special_tokens=True): return self.ocr_preprocessor.batch_decode(generated_ids, skip_special_tokens=skip_special_tokens) def crop_image(self, image, bboxes): crops_for_image = [] for bbox in bboxes: x1, y1, x2, y2 = bbox # fix the bounding box in case it is out of bounds or too small x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) x1, y1, x2, y2 = min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2) # just incase x1, y1 = max(0, x1), max(0, y1) x1, y1 = min(image.shape[1], x1), min(image.shape[0], y1) x2, y2 = max(0, x2), max(0, y2) x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) if x2 - x1 < 10: if image.shape[1] - x1 > 10: x2 = x1 + 10 else: x1 = x2 - 10 if y2 - y1 < 10: if image.shape[0] - y1 > 10: y2 = y1 + 10 else: y1 = y2 - 10 crop = image[y1:y2, x1:x2] crops_for_image.append(crop) return crops_for_image def _get_indices_of_characters_to_keep(self, batch_scores, batch_labels, batch_bboxes, character_detection_threshold): indices_of_characters_to_keep = [] for scores, labels, _ in zip(batch_scores, batch_labels, batch_bboxes): indices = torch.where((labels == 0) & (scores > character_detection_threshold))[0] indices_of_characters_to_keep.append(indices) return indices_of_characters_to_keep def _get_indices_of_panels_to_keep(self, batch_scores, batch_labels, batch_bboxes, panel_detection_threshold): indices_of_panels_to_keep = [] for scores, labels, bboxes in zip(batch_scores, batch_labels, batch_bboxes): indices = torch.where(labels == 2)[0] bboxes = bboxes[indices] scores = scores[indices] labels = labels[indices] if len(indices) == 0: indices_of_panels_to_keep.append([]) continue scores, labels, indices, bboxes = zip(*sorted(zip(scores, labels, indices, bboxes), reverse=True)) panels_to_keep = [] union_of_panels_so_far = box(0, 0, 0, 0) for ps, pb, pl, pi in zip(scores, bboxes, labels, indices): panel_polygon = box(pb[0], pb[1], pb[2], pb[3]) if ps < panel_detection_threshold: continue if union_of_panels_so_far.intersection(panel_polygon).area / panel_polygon.area > 0.5: continue panels_to_keep.append((ps, pl, pb, pi)) union_of_panels_so_far = union_of_panels_so_far.union(panel_polygon) indices_of_panels_to_keep.append([p[3].item() for p in panels_to_keep]) return indices_of_panels_to_keep def _get_indices_of_texts_to_keep(self, batch_scores, batch_labels, batch_bboxes, text_detection_threshold): indices_of_texts_to_keep = [] for scores, labels, bboxes in zip(batch_scores, batch_labels, batch_bboxes): indices = torch.where((labels == 1) & (scores > text_detection_threshold))[0] bboxes = bboxes[indices] scores = scores[indices] labels = labels[indices] if len(indices) == 0: indices_of_texts_to_keep.append([]) continue scores, labels, indices, bboxes = zip(*sorted(zip(scores, labels, indices, bboxes), reverse=True)) texts_to_keep = [] texts_to_keep_as_shapely_objects = [] for ts, tb, tl, ti in zip(scores, bboxes, labels, indices): text_polygon = box(tb[0], tb[1], tb[2], tb[3]) should_append = True for t in texts_to_keep_as_shapely_objects: if t.intersection(text_polygon).area / t.union(text_polygon).area > 0.5: should_append = False break if should_append: texts_to_keep.append((ts, tl, tb, ti)) texts_to_keep_as_shapely_objects.append(text_polygon) indices_of_texts_to_keep.append([t[3].item() for t in texts_to_keep]) return indices_of_texts_to_keep def _convert_annotations_to_coco_format(self, annotations): if annotations is None: return None self._verify_annotations_are_in_correct_format(annotations) coco_annotations = [] for annotation in annotations: coco_annotation = { "image_id": annotation["image_id"], "annotations": [], } for bbox, label in zip(annotation["bboxes_as_x1y1x2y2"], annotation["labels"]): coco_annotation["annotations"].append({ "bbox": x1y1x2y2_to_xywh(bbox), "category_id": label, "area": (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), }) coco_annotations.append(coco_annotation) return coco_annotations def _verify_annotations_are_in_correct_format(self, annotations): error_msg = """ Annotations must be in the following format: [ { "image_id": 0, "bboxes_as_x1y1x2y2": [[0, 0, 10, 10], [10, 10, 20, 20], [20, 20, 30, 30]], "labels": [0, 1, 2], }, ... ] Labels: 0 for characters, 1 for text, 2 for panels. """ if annotations is None: return if not isinstance(annotations, List) and not isinstance(annotations, tuple): raise ValueError( f"{error_msg} Expected a List/Tuple, found {type(annotations)}." ) if len(annotations) == 0: return if not isinstance(annotations[0], dict): raise ValueError( f"{error_msg} Expected a List[Dict], found {type(annotations[0])}." ) if "image_id" not in annotations[0]: raise ValueError( f"{error_msg} Dict must contain 'image_id'." ) if "bboxes_as_x1y1x2y2" not in annotations[0]: raise ValueError( f"{error_msg} Dict must contain 'bboxes_as_x1y1x2y2'." ) if "labels" not in annotations[0]: raise ValueError( f"{error_msg} Dict must contain 'labels'." )