fix segment

PVBM/DiscSegmenter.py

@@ -1,9 +1,10 @@
+import os
+import requests
 import onnxruntime as ort
 import numpy as np
 import PIL
 from PIL import Image
 from torchvision import transforms
-import os
 import cv2
 
 class DiscSegmenter:
@@ -14,6 +15,22 @@ def __init__(self):
         self.img_size = 512
         script_dir = os.path.dirname(os.path.abspath(__file__))
         self.model_path = os.path.join(script_dir, "lunetv2_odc.onnx")
+        self.download_model()
+
+    def download_model(self):
+        """Download the ONNX model if it does not exist."""
+        model_url = 'https://github.com/aim-lab/PVBM/raw/main/PVBM/lunetv2_odc.onnx'
+        print(f"Model path: {self.model_path}")
+        if not os.path.exists(self.model_path):
+            print(f"Downloading model from {model_url}...")
+            response = requests.get(model_url, stream=True)
+            response.raise_for_status()
+            with open(self.model_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print(f'Model downloaded to {self.model_path}')
+        else:
+            print("Model already exists, skipping download.")
 
     def find_biggest_contour(self, contours):
         """

pvbm.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pvbm
-Version: 2.9.3
+Version: 2.9.4
 Summary: Python Vasculature Biomarker toolbox
 Home-page: https://github.com/aim-lab/PVBM
 Author: Jonathan Fhima, Yevgeniy Men

setup.py

@@ -1,24 +1,5 @@
-import os
-import requests
 from setuptools import setup, find_packages
-from setuptools.command.install import install
 
-class PostInstallCommand(install):
-    """Post-installation for installation mode."""
-    def run(self):
-        install.run(self)
-        self.download_model()
-
-    def download_model(self):
-        model_url = 'https://github.com/aim-lab/PVBM/raw/main/PVBM/lunetv2_odc.onnx'
-        model_path = os.path.join(os.path.dirname(__file__), 'PVBM', 'lunetv2_odc.onnx')
-        if not os.path.exists(model_path):
-            response = requests.get(model_url, stream=True)
-            response.raise_for_status()
-            with open(model_path, 'wb') as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
-            print(f'Model downloaded to {model_path}')
 
 def read_readme(file_name):
     with open(file_name, "r", encoding="utf-8") as file:
@@ -28,7 +9,7 @@ def read_readme(file_name):
 
 setup(
     name='pvbm',
-    version='2.9.3',
+    version='2.9.5',
     packages=find_packages(exclude=['PVBM/lunetv2_odc.onnx']),
     install_requires=[
         "numpy",
@@ -46,7 +27,4 @@ def read_readme(file_name):
     long_description=long_description,
     long_description_content_type="text/markdown",
     url='https://github.com/aim-lab/PVBM',
-    cmdclass={
-        'install': PostInstallCommand,
-    },
 )