From 21f0b77948fbd89986f3a2e4c59b52358e1c3d76 Mon Sep 17 00:00:00 2001
From: PortfolioAI <135471798+CharlesCNorton@users.noreply.github.com>
Date: Wed, 31 Jul 2024 14:31:25 -0400
Subject: [PATCH] feat: Add model path validation and error handling (v0.7.0)

- Validate model path in `init_model`
- Check model download in `download_model`
- Enhance exception handling
---
 setup.py       |  2 +-
 yoflo/yoflo.py | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 4f8a244..8066e08 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='yoflo',
-    version='0.6.0',
+    version='0.7.0',
     packages=find_packages(),
     include_package_data=True,
     install_requires=[
diff --git a/yoflo/yoflo.py b/yoflo/yoflo.py
index fd5d2c1..caa2105 100644
--- a/yoflo/yoflo.py
+++ b/yoflo/yoflo.py
@@ -260,15 +260,26 @@ def download_model(self):
             snapshot_download(
                 repo_id="microsoft/Florence-2-base-ft", local_dir=local_model_dir
             )
-
+            if not os.path.exists(local_model_dir):
+                logging.error(
+                    f"Model download failed, directory {os.path.abspath(local_model_dir)} does not exist."
+                )
+                return False
+            if not os.path.isdir(local_model_dir):
+                logging.error(
+                    f"Model download failed, path {os.path.abspath(local_model_dir)} is not a directory."
+                )
+                return False
             logging.info(
                 f"Model and associated files downloaded and initialized at {os.path.abspath(local_model_dir)}"
             )
             self.init_model(local_model_dir)
+            return True
         except OSError as e:
             logging.error(f"OS error during model download: {e}")
         except Exception as e:
             logging.error(f"Error downloading model: {e}")
+        return False
 
     def start_webcam_detection(self):
         """Start separate threads for each specified webcam or RTSP stream."""
@@ -684,7 +695,8 @@ def main():
                 rtsp_urls=rtsp_urls,
                 record=args.record
             )
-            yo_flo.download_model()
+            if not yo_flo.download_model():
+                return
         else:
             if not os.path.exists(args.model_path):
                 logging.error(f"Model path {args.model_path} does not exist.")