Update Demo, README, and model with local tweaks

Demo.ipynb

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

README.md

@@ -130,8 +130,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_page_element, reformat_for_plotting
+from nemotron_page_elements_v3.model import define_model
+from nemotron_page_elements_v3.utils import plot_sample, postprocess_preds_page_element, reformat_for_plotting
 
 # Load image
 path = "./example.png"

nemotron_page_elements_v3/model.py

@@ -36,9 +36,9 @@ def define_model(config_name: str = "page_element_v3", verbose: bool = True) ->
     if verbose:
         print(" -> Loading weights from", config.ckpt)
 
-    # 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)
+    # 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)
 
     model = YoloXWrapper(model, config)