Introduction facilities for building a simple wheel for the nemotron-page-elements-v3 project

.gitignore

@@ -0,0 +1,8 @@
+dist/
+build/
+*.egg-info/
+*.pyc
+*.pyo
+*.pyd
+*.pyw
+*.pyz

MANIFEST.in

@@ -0,0 +1,5 @@
+include README.md
+include THIRD_PARTY_NOTICES.md
+recursive-include nemotron_page_elements_v3
+
+

model.py → nemotron_page_elements_v3/model.py

@@ -4,12 +4,13 @@ import os
 import sys
 import torch
 import importlib
+import importlib.resources
 import numpy as np
 import numpy.typing as npt
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import Dict, List, Tuple, Union
-from yolox.boxes import postprocess
+from nemotron_page_elements_v3.yolox.boxes import postprocess
 
 
 def define_model(config_name: str = "page_element_v3", verbose: bool = True) -> nn.Module:
@@ -24,8 +25,9 @@ def define_model(config_name: str = "page_element_v3", verbose: bool = True) ->
         torch.nn.Module: The initialized YOLOX model.
     """
     # Load model from exp_file
-    sys.path.append(os.path.dirname(config_name))
-    exp_module = importlib.import_module(os.path.basename(config_name).split(".")[0])
+    # page_element_v3.py is in the same directory as model.py
+    sys.path.append(os.path.dirname(__file__))
+    exp_module = importlib.import_module("page_element_v3")
 
     config = exp_module.Exp()
     model = config.get_model()
@@ -34,7 +36,9 @@ def define_model(config_name: str = "page_element_v3", verbose: bool = True) ->
     if verbose:
         print(" -> Loading weights from", config.ckpt)
 
-    ckpt = torch.load(config.ckpt, map_location="cpu", weights_only=False)
+    # Use importlib.resources to locate 'weights.pth' inside the module's directory (nmtron_page_elements_v3)
+    with importlib.resources.path("nemotron_page_elements_v3", "weights.pth") as weights_path:
+        ckpt = torch.load(str(weights_path), map_location="cpu", weights_only=False)
     model.load_state_dict(ckpt["model"], strict=True)
 
     model = YoloXWrapper(model, config)

page_element_v3.py → nemotron_page_elements_v3/page_element_v3.py

@@ -63,7 +63,9 @@ class Exp:
         Returns:
             nn.Module: The YOLOX model with configured parameters.
         """
-        from yolox import YOLOX, YOLOPAFPN, YOLOXHead
+        from nemotron_page_elements_v3.yolox.yolox import YOLOX
+        from nemotron_page_elements_v3.yolox.yolo_pafpn import YOLOPAFPN
+        from nemotron_page_elements_v3.yolox.yolo_head import YOLOXHead
 
         # Build model
         if getattr(self, "model", None) is None:

nemotron_page_elements_v3/post_processing/__init__.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Post-processing utilities for page element predictions.
+
+This module provides utilities for advanced post-processing of page element
+detection results, including box expansion, matching with titles, and
+weighted box fusion.
+"""
+
+# Import from page_elt_pp
+from .page_elt_pp import (
+    expand_boxes,
+    merge_boxes,
+    bb_iou_array,
+    match_with_title,
+    match_boxes_with_title,
+)
+
+# Import from text_pp
+from .text_pp import (
+    get_overlaps,
+    get_distances,
+    find_titles,
+    postprocess_included,
+)
+
+# Import from wbf
+from .wbf import (
+    weighted_boxes_fusion,
+    prefilter_boxes,
+    merge_labels,
+    get_weighted_box,
+    get_biggest_box,
+    find_matching_box_fast,
+)
+
+__all__ = [
+    # page_elt_pp
+    "expand_boxes",
+    "merge_boxes",
+    "bb_iou_array",
+    "match_with_title",
+    "match_boxes_with_title",
+    # text_pp
+    "get_overlaps",
+    "get_distances",
+    "find_titles",
+    "postprocess_included",
+    # wbf
+    "weighted_boxes_fusion",
+    "prefilter_boxes",
+    "merge_labels",
+    "get_weighted_box",
+    "get_biggest_box",
+    "find_matching_box_fast",
+]
+
+

pyproject.toml

@@ -0,0 +1,87 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nemotron-page-elements-v3"
+version = "3.0.0"
+description = "NVIDIA Nemotron Page Elements v3: A specialized YOLOX-based object detection model for identifying tables, charts, infographics, titles, headers/footers, and text in document pages"
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "NVIDIA Open Model License"}
+authors = [
+    {name = "NVIDIA Corporation", email = "boli@nvidia.com"},
+    {name = "Theo Viel", email = "tviel@nvidia.com"},
+    {name = "Bo Liu", email = "boli@nvidia.com"},
+]
+maintainers = [
+    {name = "Theo Viel", email = "tviel@nvidia.com"},
+    {name = "Bo Liu", email = "boli@nvidia.com"},
+]
+keywords = [
+    "object-detection",
+    "document-understanding",
+    "pdf-extraction",
+    "yolox",
+    "page-layout",
+    "nvidia",
+    "nemotron",
+    "deep-learning",
+    "computer-vision",
+]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Image Recognition",
+    "License :: Other/Proprietary License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+
+dependencies = [
+    "numpy",
+    "torch>=1.7",
+    "opencv-python",
+    "loguru",
+    "tqdm",
+    "torchvision",
+    "thop",
+    "ninja",
+    "tabulate",
+    "psutil",
+    "tensorboard",
+    "pycocotools>=2.0.2",
+    "onnx>=1.13.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "black>=22.0",
+    "flake8>=4.0",
+    "mypy>=0.950",
+]
+
+[project.urls]
+Homepage = "https://huggingface.co/nvidia/nemotron-page-elements-v3"
+Documentation = "https://huggingface.co/nvidia/nemotron-page-elements-v3"
+Repository = "https://huggingface.co/nvidia/nemotron-page-elements-v3"
+"Bug Tracker" = "https://huggingface.co/nvidia/nemotron-page-elements-v3/discussions"
+"Model Card" = "https://huggingface.co/nvidia/nemotron-page-elements-v3"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["nemotron_page_elements_v3*"]
+
+[tool.setuptools.package-data]
+"nemotron_page_elements_v3" = ["*.json", "*.pth", "*.png"]
+

yolox/__init__.py → setup.py

@@ -1,11 +1,17 @@
-#!/usr/bin/env python3
-# -*- coding:utf-8 -*-
-
+#!/usr/bin/env python
 # SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Copyright (c) Megvii Inc. All rights reserved.
+"""
+Setup script for backward compatibility.
+This project uses pyproject.toml for configuration (PEP 621).
+"""
+
+from setuptools import setup
+
+# Configuration is in pyproject.toml
+setup()
+
+
+
 
-from .yolo_head import YOLOXHead
-from .yolo_pafpn import YOLOPAFPN
-from .yolox import YOLOX