pip-install2

.gitignore

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

Demo.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:777844f15f935804cd331020a2338c7db3670a96b0dc6d3430b2602cde666e68
-size 989346
+oid sha256:4c6af3fabb5d255609ad3954c451015dbf565d3855ae54440f99429ba6978f55
+size 878101

MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include THIRD_PARTY_NOTICES.md
+recursive-include nemotron_graphic_elements_v1

README.md

@@ -134,7 +134,12 @@ git clone https://huggingface.co/nvidia/nemotron-graphic-elements-v1
 ```
 git clone [email protected]:nvidia/nemotron-graphic-elements-v1
 ```
-
+Optional:
+This can be installed as a package using pip
+```
+cd nemotron-graphic-elements-v3
+pip install -e .
+```
 2. Run the model using the following code:
 
 ```
@@ -143,8 +148,8 @@ import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 
-from model import define_model
-from utils import plot_sample, postprocess_preds_graphic_element, reformat_for_plotting
+from nemotron_graphic_elements_v1.model import define_model
+from nemotron_graphic_elements_v1.utils import plot_sample, postprocess_preds_graphic_element, reformat_for_plotting
 
 # Load image
 path = "./example.png"

nemotron_graphic_elements_v1/__init__.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Nemotron Graphic Elements v1
+
+A specialized object detection system designed to identify and extract key elements 
+from charts and graphs. Based on YOLOX architecture.
+"""
+
+__version__ = "1.0.0"
+
+from .model import define_model
+from .utils import (
+    plot_sample,
+    postprocess_preds_graphic_element,
+    reformat_for_plotting,
+    reorder_boxes,
+    COLORS,
+)
+from .graphic_element_v1 import Exp
+
+__all__ = [
+    "define_model",
+    "Exp",
+    "plot_sample",
+    "postprocess_preds_graphic_element",
+    "reformat_for_plotting",
+    "reorder_boxes",
+    "COLORS",
+]
+

graphic_element_v1.py → nemotron_graphic_elements_v1/graphic_element_v1.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 import torch
 import torch.nn as nn
 from typing import List, Tuple
@@ -18,7 +19,9 @@ class Exp:
     def __init__(self) -> None:
         """Initialize the configuration with default parameters."""
         self.name: str = "graphic-element-v1"
-        self.ckpt: str = "weights.pth"
+        # Use package directory for weights path
+        package_dir = os.path.dirname(os.path.abspath(__file__))
+        self.ckpt: str = os.path.join(package_dir, "weights.pth")
         self.device: str = "cuda:0" if torch.cuda.is_available() else "cpu"
 
         # YOLOX architecture parameters
@@ -61,7 +64,7 @@ class Exp:
         Returns:
             nn.Module: The YOLOX model with configured parameters.
         """
-        from yolox import YOLOX, YOLOPAFPN, YOLOXHead
+        from .yolox import YOLOX, YOLOPAFPN, YOLOXHead
 
         # Build model
         if getattr(self, "model", None) is None:

model.py → nemotron_graphic_elements_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 = "graphic_element_v1", verbose: bool = True) -> nn.Module:
@@ -24,11 +24,10 @@ def define_model(config_name: str = "graphic_element_v1", verbose: bool = True)
     Returns:
         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])
-
-    config = exp_module.Exp()
+    # Import the config class
+    from .graphic_element_v1 import Exp
+    
+    config = Exp()
     model = config.get_model()
 
     # Load weights

pyproject.toml

@@ -0,0 +1,54 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nemotron-graphic-elements-v1"
+version = "1.0.0"
+description = "NVIDIA Nemotron Graphic Elements v1 - A specialized object detection system for identifying and extracting key elements from charts and graphs"
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "NVIDIA Open Model License"}
+authors = [
+    {name = "NVIDIA Corporation", email = "[email protected]"}
+]
+keywords = ["object-detection", "yolox", "charts", "graphics", "pdf-ingestion", "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",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+]
+
+dependencies = [
+    "torch>=2.0.0",
+    "numpy>=1.21.0",
+    "matplotlib>=3.5.0",
+    "pandas>=1.3.0",
+    "Pillow>=9.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "jupyter>=1.0.0",
+    "notebook>=6.4.0",
+]
+
+[project.urls]
+Homepage = "https://huggingface.co/nvidia/nemotron-graphic-elements-v1"
+Documentation = "https://huggingface.co/nvidia/nemotron-graphic-elements-v1"
+Repository = "https://huggingface.co/nvidia/nemotron-graphic-elements-v1"
+"Bug Reports" = "https://app.intigriti.com/programs/nvidia/nvidiavdp/detail"
+
+[tool.setuptools]
+packages = ["nemotron_graphic_elements_v1", "nemotron_graphic_elements_v1.yolox", "nemotron_graphic_elements_v1.post_processing"]
+
+[tool.setuptools.package-data]
+"nemotron_graphic_elements_v1" = ["*.pth", "*.json", "*.png"]
+

setup.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from setuptools import setup, find_packages
+
+# For compatibility with older tools, but pyproject.toml is the primary configuration
+setup(
+    packages=find_packages(include=['nemotron_graphic_elements_v1', 'nemotron_graphic_elements_v1.*']),
+)
+