Update modelling_magi.py

modelling_magi.py

@@ -242,12 +242,15 @@ class MagiModel(PreTrainedModel):
                 file.write(transript)
         return transript
     
-    def get_text_character_affinity_matrices_given_annotations(
+    def get_affinity_matrices_given_annotations(
             self, images, annotations, move_to_device_fn=None, apply_sigmoid=True
     ):
         assert not self.config.disable_detections
         move_to_device_fn = self.move_to_device if move_to_device_fn is None else move_to_device_fn
 
+        character_bboxes_in_batch = [[bbox for bbox, label in zip(a["bboxes_as_x1y1x2y2"], a["labels"]) if label == 0] for a in annotations]
+        crop_embeddings_for_batch = self.predict_crop_embeddings(images, character_bboxes_in_batch, move_to_device_fn)
+
         inputs_to_detection_transformer = self.processor.preprocess_inputs_for_detection(images, annotations)
         inputs_to_detection_transformer = move_to_device_fn(inputs_to_detection_transformer)
         processed_targets = inputs_to_detection_transformer.pop("labels")
@@ -255,6 +258,7 @@ class MagiModel(PreTrainedModel):
         detection_transformer_output = self._get_detection_transformer_output(**inputs_to_detection_transformer)
         predicted_obj_tokens_for_batch = self._get_predicted_obj_tokens(detection_transformer_output)
         predicted_t2c_tokens_for_batch = self._get_predicted_t2c_tokens(detection_transformer_output)
+        predicted_c2c_tokens_for_batch = self._get_predicted_c2c_tokens(detection_transformer_output)
 
         predicted_class_scores, predicted_bboxes = self._get_predicted_bboxes_and_classes(detection_transformer_output)
         matching_dict = {
@@ -266,6 +270,7 @@ class MagiModel(PreTrainedModel):
         matched_char_obj_tokens_for_batch = []
         matched_text_obj_tokens_for_batch = []
         t2c_tokens_for_batch = []
+        c2c_tokens_for_batch = []
 
         text_bboxes_for_batch = []
         character_bboxes_for_batch = []
@@ -288,6 +293,7 @@ class MagiModel(PreTrainedModel):
             matched_char_obj_tokens_for_batch.append(predicted_obj_tokens_for_batch[j][predicted_char_indices])
             matched_text_obj_tokens_for_batch.append(predicted_obj_tokens_for_batch[j][predicted_text_indices])
             t2c_tokens_for_batch.append(predicted_t2c_tokens_for_batch[j])
+            c2c_tokens_for_batch.append(predicted_c2c_tokens_for_batch[j])
         
         text_character_affinity_matrices = self._get_text_character_affinity_matrices(
             character_obj_tokens_for_batch=matched_char_obj_tokens_for_batch,
@@ -296,8 +302,16 @@ class MagiModel(PreTrainedModel):
             apply_sigmoid=apply_sigmoid,
         )
 
+        character_character_affinity_matrices = self._get_character_character_affinity_matrices(
+            character_obj_tokens_for_batch=matched_char_obj_tokens_for_batch,
+            crop_embeddings_for_batch=crop_embeddings_for_batch,
+            c2c_tokens_for_batch=c2c_tokens_for_batch,
+            apply_sigmoid=apply_sigmoid,
+        )
+
         return {
             "text_character_affinity_matrices": text_character_affinity_matrices,
+            "character_character_affinity_matrices": character_character_affinity_matrices,
             "text_bboxes_for_batch": text_bboxes_for_batch,
             "character_bboxes_for_batch": character_bboxes_for_batch,
         }