enhance logging for model loading and pipeline processing; update log levels and add detailed error messages

siwatsystem/pympta.py

@@ -6,19 +6,27 @@ import cv2
 import requests
 import zipfile
 import shutil
+import traceback
 from ultralytics import YOLO
 from urllib.parse import urlparse
 
+# Create a logger specifically for this module
+logger = logging.getLogger("detector_worker.pympta")
+
 def load_pipeline_node(node_config: dict, mpta_dir: str) -> dict:
     # Recursively load a model node from configuration.
     model_path = os.path.join(mpta_dir, node_config["modelFile"])
     if not os.path.exists(model_path):
-        logging.error(f"Model file {model_path} not found.")
+        logger.error(f"Model file {model_path} not found. Current directory: {os.getcwd()}")
+        logger.error(f"Directory content: {os.listdir(os.path.dirname(model_path))}")
         raise FileNotFoundError(f"Model file {model_path} not found.")
-    logging.info(f"Loading model for node {node_config['modelId']} from {model_path}")
+    logger.info(f"Loading model for node {node_config['modelId']} from {model_path}")
     model = YOLO(model_path)
     if torch.cuda.is_available():
+        logger.info(f"CUDA available. Moving model {node_config['modelId']} to GPU")
         model.to("cuda")
+    else:
+        logger.info(f"CUDA not available. Using CPU for model {node_config['modelId']}")
     node = {
         "modelId": node_config["modelId"],
         "modelFile": node_config["modelFile"],
@@ -28,11 +36,14 @@ def load_pipeline_node(node_config: dict, mpta_dir: str) -> dict:
         "model": model,
         "branches": []
     }
+    logger.debug(f"Configured node {node_config['modelId']} with trigger classes: {node['triggerClasses']}")
     for child in node_config.get("branches", []):
+        logger.debug(f"Loading branch for parent node {node_config['modelId']}")
         node["branches"].append(load_pipeline_node(child, mpta_dir))
     return node
 
 def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict:
+    logger.info(f"Attempting to load pipeline from {zip_source} to {target_dir}")
     os.makedirs(target_dir, exist_ok=True)
     zip_path = os.path.join(target_dir, "pipeline.mpta")
     
@@ -40,51 +51,121 @@ def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict:
     parsed = urlparse(zip_source)
     if parsed.scheme in ("", "file"):
         local_path = parsed.path if parsed.scheme == "file" else zip_source
+        logger.debug(f"Checking if local file exists: {local_path}")
         if os.path.exists(local_path):
             try:
                 shutil.copy(local_path, zip_path)
-                logging.info(f"Copied local .mpta file from {local_path} to {zip_path}")
+                logger.info(f"Copied local .mpta file from {local_path} to {zip_path}")
             except Exception as e:
-                logging.error(f"Failed to copy local .mpta file from {local_path}: {e}")
+                logger.error(f"Failed to copy local .mpta file from {local_path}: {str(e)}", exc_info=True)
                 return None
         else:
-            logging.error(f"Local file {local_path} does not exist.")
+            logger.error(f"Local file {local_path} does not exist. Current directory: {os.getcwd()}")
+            # List all subdirectories of models directory to help debugging
+            if os.path.exists("models"):
+                logger.error(f"Content of models directory: {os.listdir('models')}")
+                for root, dirs, files in os.walk("models"):
+                    logger.error(f"Directory {root} contains subdirs: {dirs} and files: {files}")
+            else:
+                logger.error("The models directory doesn't exist")
             return None
     else:
-        logging.error("HTTP download functionality has been moved. Use a local file path here.")
+        logger.error(f"HTTP download functionality has been moved. Use a local file path here. Received: {zip_source}")
         return None
 
     try:
+        if not os.path.exists(zip_path):
+            logger.error(f"Zip file not found at expected location: {zip_path}")
+            return None
+            
+        logger.debug(f"Extracting .mpta file from {zip_path} to {target_dir}")
+        # Extract contents and track the directories created
+        extracted_dirs = []
         with zipfile.ZipFile(zip_path, "r") as zip_ref:
+            file_list = zip_ref.namelist()
+            logger.debug(f"Files in .mpta archive: {file_list}")
+            
+            # Extract and track the top-level directories
+            for file_path in file_list:
+                parts = file_path.split('/')
+                if len(parts) > 1:
+                    top_dir = parts[0]
+                    if top_dir and top_dir not in extracted_dirs:
+                        extracted_dirs.append(top_dir)
+            
+            # Now extract the files
             zip_ref.extractall(target_dir)
-        logging.info(f"Extracted .mpta file to {target_dir}")
+            
+        logger.info(f"Successfully extracted .mpta file to {target_dir}")
+        logger.debug(f"Extracted directories: {extracted_dirs}")
+        
+        # Check what was actually created after extraction
+        actual_dirs = [d for d in os.listdir(target_dir) if os.path.isdir(os.path.join(target_dir, d))]
+        logger.debug(f"Actual directories created: {actual_dirs}")
+    except zipfile.BadZipFile as e:
+        logger.error(f"Bad zip file {zip_path}: {str(e)}", exc_info=True)
+        return None
     except Exception as e:
-        logging.error(f"Failed to extract .mpta file: {e}")
+        logger.error(f"Failed to extract .mpta file {zip_path}: {str(e)}", exc_info=True)
         return None
     finally:
         if os.path.exists(zip_path):
             os.remove(zip_path)
+            logger.debug(f"Removed temporary zip file: {zip_path}")
+
+    # Use the first extracted directory if it exists, otherwise use the expected name
     pipeline_name = os.path.basename(zip_source)
     pipeline_name = os.path.splitext(pipeline_name)[0]
-    mpta_dir = os.path.join(target_dir, pipeline_name)
+    
+    # Find the directory with pipeline.json
+    mpta_dir = None
+    # First try the expected directory name
+    expected_dir = os.path.join(target_dir, pipeline_name)
+    if os.path.exists(expected_dir) and os.path.exists(os.path.join(expected_dir, "pipeline.json")):
+        mpta_dir = expected_dir
+        logger.debug(f"Found pipeline.json in the expected directory: {mpta_dir}")
+    else:
+        # Look through all subdirectories for pipeline.json
+        for subdir in actual_dirs:
+            potential_dir = os.path.join(target_dir, subdir)
+            if os.path.exists(os.path.join(potential_dir, "pipeline.json")):
+                mpta_dir = potential_dir
+                logger.info(f"Found pipeline.json in directory: {mpta_dir} (different from expected: {expected_dir})")
+                break
+    
+    if not mpta_dir:
+        logger.error(f"Could not find pipeline.json in any extracted directory. Directory content: {os.listdir(target_dir)}")
+        return None
+        
     pipeline_json_path = os.path.join(mpta_dir, "pipeline.json")
     if not os.path.exists(pipeline_json_path):
-        logging.error("pipeline.json not found in the .mpta file")
+        logger.error(f"pipeline.json not found in the .mpta file. Files in directory: {os.listdir(mpta_dir)}")
         return None
 
     try:
         with open(pipeline_json_path, "r") as f:
             pipeline_config = json.load(f)
+        logger.info(f"Successfully loaded pipeline configuration from {pipeline_json_path}")
+        logger.debug(f"Pipeline config: {json.dumps(pipeline_config, indent=2)}")
         return load_pipeline_node(pipeline_config["pipeline"], mpta_dir)
+    except json.JSONDecodeError as e:
+        logger.error(f"Error parsing pipeline.json: {str(e)}", exc_info=True)
+        return None
+    except KeyError as e:
+        logger.error(f"Missing key in pipeline.json: {str(e)}", exc_info=True)
+        return None
     except Exception as e:
-        logging.error(f"Error loading pipeline.json: {e}")
+        logger.error(f"Error loading pipeline.json: {str(e)}", exc_info=True)
         return None
 
-def run_pipeline(frame, node: dict, return_bbox: bool = False):
+def run_pipeline(frame, node: dict, return_bbox: bool = False, is_last_stage: bool = True):
     """
     Processes the frame with the given pipeline node. When return_bbox is True,
     the function returns a tuple (detection, bbox) where bbox is (x1,y1,x2,y2)
     for drawing. Otherwise, returns only the detection.
+    
+    The is_last_stage parameter controls whether this node is considered the last
+    in the pipeline chain. Only the last stage will return detection results.
     """
     try:
         # Check model type and use appropriate method
@@ -92,7 +173,7 @@ def run_pipeline(frame, node: dict, return_bbox: bool = False):
         
         if model_task == "classify":
             # Classification models need to use predict() instead of track()
-            logging.debug(f"Running classification model: {node.get('modelId')}")
+            logger.debug(f"Running classification model: {node.get('modelId')}")
             results = node["model"].predict(frame, stream=False)
             detection = None
             best_box = None
@@ -109,18 +190,32 @@ def run_pipeline(frame, node: dict, return_bbox: bool = False):
                         "confidence": conf,
                         "id": None  # Classification doesn't have tracking IDs
                     }
+                    logger.debug(f"Classification detection: {detection}")
+                else:
+                    logger.debug(f"Empty classification results for model {node.get('modelId')}")
             
             # Classification doesn't produce bounding boxes
             bbox = None
             
         else:
             # Detection/segmentation models use tracking
-            logging.debug(f"Running detection/tracking model: {node.get('modelId')}")
+            logger.debug(f"Running detection/tracking model: {node.get('modelId')}")
             results = node["model"].track(frame, stream=False, persist=True)
             detection = None
             best_box = None
             max_conf = -1
 
+            # Log raw detection count
+            detection_count = 0
+            for r in results:
+                if hasattr(r.boxes, 'cpu') and len(r.boxes.cpu()) > 0:
+                    detection_count += len(r.boxes.cpu())
+            
+            if detection_count == 0:
+                logger.debug(f"Empty detection results (no objects found) for model {node.get('modelId')}")
+            else:
+                logger.debug(f"Detection model {node.get('modelId')} found {detection_count} objects")
+
             for r in results:
                 for box in r.boxes:
                     box_cpu = box.cpu()
@@ -133,6 +228,11 @@ def run_pipeline(frame, node: dict, return_bbox: bool = False):
                             "id": box.id.item()
                         }
                         best_box = box_cpu
+            
+            if detection:
+                logger.debug(f"Best detection: {detection}")
+            else:
+                logger.debug(f"No valid detection with tracking ID for model {node.get('modelId')}")
 
             bbox = None
             # Calculate bbox if best_box exists
@@ -144,31 +244,44 @@ def run_pipeline(frame, node: dict, return_bbox: bool = False):
                 x2, y2 = min(w, x2), min(h, y2)
                 if x2 > x1 and y2 > y1:
                     bbox = (x1, y1, x2, y2)
+                    logger.debug(f"Detection bounding box: {bbox}")
                     if node.get("crop", False):
                         frame = frame[y1:y2, x1:x2]
+                        logger.debug(f"Cropped frame to {frame.shape}")
 
+        # Check if we should process branches
         if detection is not None:
             for branch in node["branches"]:
                 if detection["class"] in branch.get("triggerClasses", []):
                     min_conf = branch.get("minConfidence")
                     if min_conf is not None and detection["confidence"] < min_conf:
-                        logging.debug(f"Confidence {detection['confidence']} below threshold {min_conf} for branch {branch['modelId']}.")
+                        logger.debug(f"Confidence {detection['confidence']} below threshold {min_conf} for branch {branch['modelId']}.")
+                        break
+                    
+                    # If we have branches, this is not the last stage
+                    branch_result = run_pipeline(frame, branch, return_bbox, is_last_stage=True)
+                    
+                    # This node is no longer the last stage, so its results shouldn't be returned
+                    is_last_stage = False
+                    
+                    if branch_result is not None:
                         if return_bbox:
-                            return detection, bbox
-                        return detection
-                    res = run_pipeline(frame, branch, return_bbox)
-                    if res is not None:
-                        if return_bbox:
-                            return res
-                        return res
-            if return_bbox:
-                return detection, bbox
-            return detection
+                            return branch_result
+                        return branch_result
+                    break
+            
+            # Return this node's detection only if it's considered the last stage
+            if is_last_stage:
+                if return_bbox:
+                    return detection, bbox
+                return detection
+            
+        # No detection or not the last stage
         if return_bbox:
             return None, None
         return None
     except Exception as e:
-        logging.error(f"Error running pipeline on node {node.get('modelId')}: {e}")
+        logger.error(f"Error running pipeline on node {node.get('modelId')}: {e}")
         if return_bbox:
             return None, None
         return None