| 70 | |
| 71 | |
| 72 | class OnnxModel(DocLayoutModel): |
| 73 | def __init__(self, model_path: str): |
| 74 | self.model_path = model_path |
| 75 | |
| 76 | model = onnx.load(model_path) |
| 77 | metadata = {d.key: d.value for d in model.metadata_props} |
| 78 | self._stride = ast.literal_eval(metadata["stride"]) |
| 79 | self._names = ast.literal_eval(metadata["names"]) |
| 80 | |
| 81 | self.model = onnxruntime.InferenceSession(model.SerializeToString()) |
| 82 | |
| 83 | @staticmethod |
| 84 | def from_pretrained(): |
| 85 | pth = get_doclayout_onnx_model_path() |
| 86 | return OnnxModel(pth) |
| 87 | |
| 88 | @property |
| 89 | def stride(self): |
| 90 | return self._stride |
| 91 | |
| 92 | def resize_and_pad_image(self, image, new_shape): |
| 93 | """ |
| 94 | Resize and pad the image to the specified size, ensuring dimensions are multiples of stride. |
| 95 | |
| 96 | Parameters: |
| 97 | - image: Input image |
| 98 | - new_shape: Target size (integer or (height, width) tuple) |
| 99 | - stride: Padding alignment stride, default 32 |
| 100 | |
| 101 | Returns: |
| 102 | - Processed image |
| 103 | """ |
| 104 | if isinstance(new_shape, int): |
| 105 | new_shape = (new_shape, new_shape) |
| 106 | |
| 107 | h, w = image.shape[:2] |
| 108 | new_h, new_w = new_shape |
| 109 | |
| 110 | # Calculate scaling ratio |
| 111 | r = min(new_h / h, new_w / w) |
| 112 | resized_h, resized_w = int(round(h * r)), int(round(w * r)) |
| 113 | |
| 114 | # Resize image |
| 115 | image = cv2.resize( |
| 116 | image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR |
| 117 | ) |
| 118 | |
| 119 | # Calculate padding size and align to stride multiple |
| 120 | pad_w = (new_w - resized_w) % self.stride |
| 121 | pad_h = (new_h - resized_h) % self.stride |
| 122 | top, bottom = pad_h // 2, pad_h - pad_h // 2 |
| 123 | left, right = pad_w // 2, pad_w - pad_w // 2 |
| 124 | |
| 125 | # Add padding |
| 126 | image = cv2.copyMakeBorder( |
| 127 | image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114) |
| 128 | ) |
| 129 |
no outgoing calls