updated model path, readme and demo

Demo.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:433fa0f144f075b75de5f2fa2925a808b6b15c18eaf58f591e09f50c3e121b9c
-size 1812492
+oid sha256:1abc855769ac788f46bc709cd33405fe02c816c68ad500a75ae82444998fbb61
+size 1812701

README.md

@@ -121,7 +121,12 @@ git clone https://huggingface.co/nvidia/nemotron-page-elements-v3
 ```
 git clone [email protected]:nvidia/nemotron-page-elements-v3
 ```
-
+Optional:
+This can be installed as a package using pip
+```
+cd nemotron-page-elements-v3
+pip install -e .
+```
 2. Run the model using the following code:
 
 ```

nemotron_page_elements_v3/model.py

@@ -36,15 +36,15 @@ def define_model(config_name: str = "page_element_v3", verbose: bool = True) ->
     if verbose:
         print(" -> Loading weights from", config.ckpt)
 
-    # Locate 'weights.pth' in the same directory as this model.py file
-    weights_path = os.path.join(os.path.dirname(__file__), "weights.pth")
-    ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
-    model.load_state_dict(ckpt["model"], strict=True)
-
+    # Find package directory and load weights (nemotron_page_elements_v3)
+    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)
-
-
+    
 def resize_pad(img: torch.Tensor, size: tuple) -> torch.Tensor:
     """
     Resizes and pads an image to a given size.