Duplicate from wybxc/DocLayout-YOLO-DocStructBench-onnx

Co-authored-by: wybxc <wybxc@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,6 @@
+---
+license: apache-2.0
+base_model:
+- juliozhao/DocLayout-YOLO-DocStructBench
+---
+Converted from [juliozhao/DocLayout-YOLO-DocStructBench](https://huggingface.co/juliozhao/DocLayout-YOLO-DocStructBench)

doclayout_yolo_docstructbench_imgsz1024.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fece9af02f618b603ff7921ccec6861d13e7e1f9830e091dfb7e8ad9311e5b21
+size 75324598

inference.py

@@ -0,0 +1,174 @@
+import ast
+import onnx
+import onnxruntime as ort
+import cv2
+from huggingface_hub import hf_hub_download
+import numpy as np
+
+# Download the model from the Hugging Face Hub
+model = hf_hub_download(
+    repo_id="wybxc/DocLayout-YOLO-DocStructBench-onnx",
+    filename="doclayout_yolo_docstructbench_imgsz1024.onnx",
+)
+model = onnx.load(model)
+metadata = {prop.key: prop.value for prop in model.metadata_props}
+
+names = ast.literal_eval(metadata["names"])
+stride = ast.literal_eval(metadata["stride"])
+
+# Load the model with ONNX Runtime
+session = ort.InferenceSession(model.SerializeToString())
+
+
+def resize_and_pad_image(image, new_shape, stride=32):
+    """
+    Resize and pad the image to the specified size, ensuring dimensions are multiples of stride.
+
+    Parameters:
+    - image: Input image
+    - new_shape: Target size (integer or (height, width) tuple)
+    - stride: Padding alignment stride, default 32
+
+    Returns:
+    - Processed image
+    """
+    if isinstance(new_shape, int):
+        new_shape = (new_shape, new_shape)
+
+    h, w = image.shape[:2]
+    new_h, new_w = new_shape
+
+    # Calculate scaling ratio
+    r = min(new_h / h, new_w / w)
+    resized_h, resized_w = int(round(h * r)), int(round(w * r))
+
+    # Resize image
+    image = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR)
+
+    # Calculate padding size and align to stride multiple
+    pad_w = (new_w - resized_w) % stride
+    pad_h = (new_h - resized_h) % stride
+    top, bottom = pad_h // 2, pad_h - pad_h // 2
+    left, right = pad_w // 2, pad_w - pad_w // 2
+
+    # Add padding
+    image = cv2.copyMakeBorder(
+        image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)
+    )
+
+    return image
+
+
+def scale_boxes(img1_shape, boxes, img0_shape):
+    """
+    Rescales bounding boxes (in the format of xyxy by default) from the shape of the image they were originally
+    specified in (img1_shape) to the shape of a different image (img0_shape).
+
+    Args:
+        img1_shape (tuple): The shape of the image that the bounding boxes are for, in the format of (height, width).
+        boxes (torch.Tensor): the bounding boxes of the objects in the image, in the format of (x1, y1, x2, y2)
+        img0_shape (tuple): the shape of the target image, in the format of (height, width).
+
+    Returns:
+        boxes (torch.Tensor): The scaled bounding boxes, in the format of (x1, y1, x2, y2)
+    """
+
+    # Calculate scaling ratio
+    gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
+
+    # Calculate padding size
+    pad_x = round((img1_shape[1] - img0_shape[1] * gain) / 2 - 0.1)
+    pad_y = round((img1_shape[0] - img0_shape[0] * gain) / 2 - 0.1)
+
+    # Remove padding and scale boxes
+    boxes[..., :4] = (boxes[..., :4] - [pad_x, pad_y, pad_x, pad_y]) / gain
+    return boxes
+
+
+class YoloResult:
+    def __init__(self, boxes, names):
+        self.boxes = [YoloBox(data=d) for d in boxes]
+        self.boxes = sorted(self.boxes, key=lambda x: x.conf, reverse=True)
+        self.names = names
+
+
+class YoloBox:
+    def __init__(self, data):
+        self.xyxy = data[:4]
+        self.conf = data[-2]
+        self.cls = data[-1]
+
+
+def inference(image):
+    """
+    Run inference on the input image.
+
+    Parameters:
+    - image: Input image, HWC format and RGB order
+
+    Returns:
+    - YoloResult object containing the predicted boxes and class names
+    """
+
+    # Preprocess image
+    orig_h, orig_w = image.shape[:2]
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    pix = resize_and_pad_image(image, new_shape=int(image.shape[0] / stride) * stride)
+    pix = np.transpose(pix, (2, 0, 1))  # CHW
+    pix = np.expand_dims(pix, axis=0)  # BCHW
+    pix = pix.astype(np.float32) / 255.0  # Normalize to [0, 1]
+    new_h, new_w = pix.shape[2:]
+
+    # Run inference
+    preds = session.run(None, {"images": pix})[0]
+
+    # Postprocess predictions
+    preds = preds[preds[..., 4] > 0.25]
+    preds[..., :4] = scale_boxes((new_h, new_w), preds[..., :4], (orig_h, orig_w))
+    return YoloResult(boxes=preds, names=names)
+
+
+if __name__ == "__main__":
+    import sys
+    import matplotlib
+    import matplotlib.pyplot as plt
+    import matplotlib.colors as colors
+
+    image = sys.argv[1]
+    image = cv2.imread(image)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    layout = inference(image)
+
+    bitmap = np.ones(image.shape[:2], dtype=np.uint8)
+    h, w = bitmap.shape
+    vcls = ["abandon", "figure", "table", "isolate_formula", "formula_caption"]
+    for i, d in enumerate(layout.boxes):
+        x0, y0, x1, y1 = d.xyxy.squeeze()
+        x0, y0, x1, y1 = (
+            np.clip(int(x0 - 1), 0, w - 1),
+            np.clip(int(h - y1 - 1), 0, h - 1),
+            np.clip(int(x1 + 1), 0, w - 1),
+            np.clip(int(h - y0 + 1), 0, h - 1),
+        )
+        if layout.names[int(d.cls)] in vcls:
+            bitmap[y0:y1, x0:x1] = 0
+        else:
+            bitmap[y0:y1, x0:x1] = i + 2
+    bitmap = bitmap[::-1, :]
+
+    # map bitmap to color
+    colormap = matplotlib.colormaps["Pastel1"]
+    norm = colors.Normalize(vmin=bitmap.min(), vmax=bitmap.max())
+    colored_bitmap = colormap(norm(bitmap))
+    colored_bitmap = (colored_bitmap[:, :, :3] * 255).astype(np.uint8)
+
+    # overlay bitmap on image
+    image_with_bitmap = cv2.multiply(image, colored_bitmap, scale=1 / 255)
+
+    # show the results
+    fig, ax = plt.subplots(1, 3, figsize=(15, 6))
+    ax[0].imshow(image)
+    ax[1].imshow(bitmap, cmap="Pastel1")
+    ax[2].imshow(image_with_bitmap)
+    plt.show()