RE_UPLOAD-REBUILD-RESTART

utils/get_features.py

@@ -0,0 +1,123 @@
+import layoutparser as lp
+from PIL import Image
+import tensorflow as tf
+import numpy as np
+import torch
+import torchvision.ops.boxes as box_ops
+from typing import List, Tuple
+from .split_image import split_image
+from .get_unique_values import get_unique_values
+
+def get_vectors(*,
+    predicted_bboxes: List[Tuple[int, int, int, int]],
+    predicted_scores: List[float],
+    predicted_labels: List[str],
+    label_names: List[str],
+    sub_images_bboxes: List[Tuple[int, int, int, int]],
+    index_start: int = 0.17,
+    index_end: int = 1,
+    weighted_jaccard_index = False):
+  bboxes_tensor: torch.Tensor = torch.tensor(predicted_bboxes)
+  labels_nonce = { value:key for key, value in zip(get_unique_values(start = index_start, end = index_end, count = len(label_names)), list(label_names)) }
+  
+  def get_vector(bbox: Tuple[int, int, int, int], region_nonce: int):
+    # bbox: Expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
+    bbox_tensor: torch.Tensor = torch.tensor([bbox])
+    [jaccard_indexes] = box_ops.box_iou(bbox_tensor, bboxes_tensor)
+    '''
+      Either get the index of bounding box with largest jaccard_index (Intersection Over Union) or
+      get the index of bounding box with largest jaccard_index (Intersection Over Union) multiplied by the score.
+      By doing this we strike a balance between accuracy and relative position.
+    '''
+    index_of_jaccard_index = jaccard_indexes.argmax() if not weighted_jaccard_index else np.multiply(jaccard_indexes, predicted_scores).argmax()
+    jaccard_index = jaccard_indexes[index_of_jaccard_index]
+    jaccard_index_bbox_label__nonce = labels_nonce[predicted_labels[index_of_jaccard_index]]
+    jaccard_index_bbox_score = predicted_scores[index_of_jaccard_index]
+    vector = region_nonce * jaccard_index * jaccard_index_bbox_label__nonce * jaccard_index_bbox_score
+    return vector.item()
+  sub_images_nonces = get_unique_values(start = index_start, end = index_end, count = len(sub_images_bboxes))
+  for sub_image_bbox, region_nonce in zip(sub_images_bboxes, sub_images_nonces):
+      yield get_vector(sub_image_bbox, region_nonce)
+
+def get_predictions(
+    image: Image.Image,
+    model: lp.Detectron2LayoutModel,
+    predictions_reducer = lambda *args: args):
+  layout_predicted = model.detect(image)
+  if len(layout_predicted) > 0:
+    predicted_bboxes = [block.coordinates for block in layout_predicted]
+    predicted_scores = [block.score for block in layout_predicted]
+    predicted_labels = [block.type for block in layout_predicted]
+    [predicted_bboxes, predicted_scores, predicted_labels] = predictions_reducer(
+        predicted_bboxes,
+        predicted_scores,
+        predicted_labels)
+    return {
+        'predicted_bboxes': predicted_bboxes,
+        'predicted_scores': predicted_scores,
+        'predicted_labels': predicted_labels,
+    }
+  else:
+    return {
+        'predicted_bboxes': [],
+        'predicted_scores': [],
+        'predicted_labels': [],
+    }
+
+def predictions_reducer(
+    predicted_bboxes: List[Tuple[int, int, int, int]],
+    predicted_scores: List[float],
+    predicted_labels: List[str]):
+  selected_indices = tf.image.non_max_suppression(
+    boxes = predicted_bboxes,
+    scores = predicted_scores ,
+    max_output_size = len(predicted_bboxes),
+    iou_threshold = 0.01)
+  return {
+      'predicted_bboxes': tf.gather(predicted_bboxes, selected_indices).numpy().tolist(), # List[List[int, int, int, int]]
+      'predicted_scores': tf.gather(predicted_scores, selected_indices).numpy().astype(float).tolist(),
+      'predicted_labels': tf.gather(predicted_labels, selected_indices).numpy().astype(str).tolist()
+    }
+
+def get_features(image: Image.Image, model: lp.Detectron2LayoutModel, label_names: List[str], width_parts = 100, height_parts = 100):
+  predictions = get_predictions(image, model)
+  reduced_predictions = predictions_reducer(**predictions)
+  sub_images_bboxes = list(split_image(np.array(image), width_parts, height_parts, result = 'bboxes'))
+
+  vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = False,
+    **predictions)
+  
+  weighted_vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = True,
+    **predictions)
+  
+  reduced_vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = False,
+    **reduced_predictions)
+  
+  reduced_weighted_vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = True,
+    **reduced_predictions)
+
+  return {
+      'predicted_bboxes': predictions['predicted_bboxes'],
+      'predicted_scores': predictions['predicted_scores'],
+      'predicted_labels': predictions['predicted_labels'],
+      'vectors': list(vectors),
+      'weighted_vectors': list(weighted_vectors),
+
+      'reduced_predicted_bboxes': reduced_predictions['predicted_bboxes'],
+      'reduced_predicted_scores': reduced_predictions['predicted_scores'],
+      'reduced_predicted_labels': reduced_predictions['predicted_labels'],
+      'reduced_vectors': list(reduced_vectors),
+      'reduced_weighted_vectors': list(reduced_weighted_vectors),
+  }