pip-install2 (#4)

- restructure so that table-structure can be installed via pip (697b9172fdff60b1635324cf971589781d14ee1c)
- Clean up repository structure and update for pip install (12ea31ec447e5d26a4a4ddda5b51bd3f6a2ae5d3)
- removed redundant utils.py (356eefc12f9211303b630a796b0cba9c45437c06)


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

.gitignore

@@ -0,0 +1,7 @@
+dist
+build
+*.egg-info
+*.egg
+*.pyc
+*.pyo
+*.pyd

Demo.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3fbac9002e8052dd97f8de19ede6a60fa119d237a08aa13e09fc294c73708489
-size 1085041
+oid sha256:002e6edf2b37d18f5eb2499fa653b8543b95964b122be8726cf214a2cf5500ba
+size 848913

MANIFEST.in

@@ -0,0 +1,6 @@
+include README.md
+include THIRD_PARTY_NOTICES.md
+recursive-include nemotron_table_structure_v1
+
+
+

README.md

@@ -134,7 +134,12 @@ git clone https://huggingface.co/nvidia/nemotron-table-structure-v1
 ```
 git clone git@hf.co:nvidia/nemotron-table-structure-v1
 ```
-
+Optional:
+This can be installed as a package using pip
+```
+cd nemotron-table-structure-v1
+pip install -e .
+```
 2. Run the model using the following code:
 
 ```

nemotron_table_structure_v1/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Nemotron Table Structure v1
+
+A specialized object detection model for table structure extraction based on YOLOX.
+"""
+
+__version__ = "1.0.0"
+
+from .model import define_model, YoloXWrapper
+from .utils import (
+    plot_sample,
+    postprocess_preds_table_structure,
+    reformat_for_plotting,
+    reorder_boxes,
+)
+
+__all__ = [
+    "define_model",
+    "YoloXWrapper",
+    "plot_sample",
+    "postprocess_preds_table_structure",
+    "reformat_for_plotting",
+    "reorder_boxes",
+]
+
+
+

model.py → nemotron_table_structure_v1/model.py

@@ -10,7 +10,7 @@ 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 .yolox.boxes import postprocess
 
 
 def define_model(config_name: str = "page_element_v3", verbose: bool = True) -> nn.Module:
@@ -25,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("table_structure_v1")
 
     config = exp_module.Exp()
     model = config.get_model()
@@ -35,8 +36,11 @@ 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)
-    model.load_state_dict(ckpt["model"], strict=True)
+    # Find package directory and load weights (nemotron_table_structure_v1)
+    package_dir = os.path.dirname(os.path.abspath(__file__))
+    weights_path = os.path.join(package_dir, "weights.pth")
+    state_dict = torch.load(weights_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(state_dict["model"], strict=True)
 
     model = YoloXWrapper(model, config)
     return model.eval().to(config.device)

nemotron_table_structure_v1/post_processing/__init__.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .table_struct_pp import (
+    assign_boxes,
+    merge_text_in_cell,
+    remove_empty_row,
+    build_markdown,
+    display_markdown,
+)
+from .wbf import weighted_boxes_fusion
+
+__all__ = [
+    "assign_boxes",
+    "merge_text_in_cell",
+    "remove_empty_row",
+    "build_markdown",
+    "display_markdown",
+    "weighted_boxes_fusion",
+]
+
+
+

table_structure_v1.py → nemotron_table_structure_v1/table_structure_v1.py

@@ -17,7 +17,7 @@ class Exp:
 
     def __init__(self) -> None:
         """Initialize the configuration with default parameters."""
-        self.name: str = "page-element-v3"
+        self.name: str = "table-structure-v1"
         self.ckpt: str = "weights.pth"
         self.device: str = "cuda:0" if torch.cuda.is_available() else "cpu"
 
@@ -56,7 +56,9 @@ class Exp:
         Returns:
             nn.Module: The YOLOX model with configured parameters.
         """
-        from yolox import YOLOX, YOLOPAFPN, YOLOXHead
+        from nemotron_table_structure_v1.yolox.yolox import YOLOX
+        from nemotron_table_structure_v1.yolox.yolo_pafpn import YOLOPAFPN
+        from nemotron_table_structure_v1.yolox.yolo_head import YOLOXHead
 
         # Build model
         if getattr(self, "model", None) is None:

nemotron_table_structure_v1/weights.pth

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

pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["setuptools>=45", "wheel", "setuptools_scm[toml]>=6.2"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nemotron-table-structure-v1"
+version = "1.0.0"
+description = "Nemotron Table Structure v1 - A specialized object detection model for table structure extraction"
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "NVIDIA Open Model License"}
+authors = [
+    {name = "NVIDIA Corporation", email = "tviel@nvidia.com"}
+]
+keywords = ["table-structure", "object-detection", "yolox", "nvidia", "nemotron", "pdf", "document-processing", "ocr"]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: Other/Proprietary License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Image Recognition",
+]
+dependencies = [
+    "torch>=2.0.0",
+    "numpy>=1.21.0",
+    "matplotlib>=3.3.0",
+    "pandas>=1.3.0",
+    "Pillow>=8.0.0",
+]
+
+[project.urls]
+Homepage = "https://huggingface.co/nvidia/nemotron-table-structure-v1"
+Repository = "https://huggingface.co/nvidia/nemotron-table-structure-v1"
+Documentation = "https://huggingface.co/nvidia/nemotron-table-structure-v1"
+"Bug Tracker" = "https://huggingface.co/nvidia/nemotron-table-structure-v1/discussions"
+
+[tool.setuptools]
+packages = ["nemotron_table_structure_v1", "nemotron_table_structure_v1.yolox", "nemotron_table_structure_v1.post_processing"]
+
+[tool.setuptools.package-data]
+"*" = ["*.pth", "config.json"]
+
+
+