endoreg-db 0.8.6.1__py3-none-any.whl → 0.8.8.0__py3-none-any.whl

endoreg_db/models/media/video/video_file_ai.py

@@ -1,21 +1,173 @@
 import logging
+from collections import Counter, defaultdict
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, List, Dict, Tuple
-from collections import defaultdict, Counter
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from safetensors import safe_open
 
-from ...utils import TEST_RUN as GLOBAL_TEST_RUN, TEST_RUN_FRAME_NUMBER as GLOBAL_N_TEST_FRAMES
 from ...metadata import ModelMeta, VideoPredictionMeta
+from ...utils import TEST_RUN as GLOBAL_TEST_RUN
+from ...utils import TEST_RUN_FRAME_NUMBER as GLOBAL_N_TEST_FRAMES
 
 if TYPE_CHECKING:
-    from .video_file import VideoFile
     from ...medical.hardware import EndoscopyProcessor
+    from .video_file import VideoFile
 
 logger = logging.getLogger(__name__)
 
 
-def _extract_text_from_video_frames(
-    video: "VideoFile", frame_fraction: float = 0.001, cap: int = 15
-) -> Optional[Dict[str, str]]:
+def _is_stub_weights_file(weights_path: Path) -> bool:
+    """Return True if the provided weights file is a known test stub."""
+
+    name_hint = weights_path.name.lower()
+    if "stub" in name_hint:
+        return True
+
+    try:
+        size_bytes = weights_path.stat().st_size
+    except OSError:
+        return False
+
+    if size_bytes < 4096:
+        try:
+            with weights_path.open("rb") as fh:
+                header = fh.read(32)
+        except OSError:
+            return False
+        return header.startswith(b"stub-weights") or not header
+
+    return False
+
+
+def _resolve_label_names(model_meta: "ModelMeta") -> List[str]:
+    """Return deterministic label ordering for the associated label set."""
+
+    labelset = model_meta.labelset
+    if not labelset:
+        return []
+
+    try:
+        return [label.name for label in labelset.get_labels_in_order()]
+    except AttributeError:
+        # Fallback in case legacy labelsets provide only the raw manager interface.
+        return [label.name for label in labelset.labels.all().order_by("name")]
+
+
+def _infer_model_type(model_meta: "ModelMeta", weights_path: Path) -> str:
+    """Best-effort detection of the backbone expected by the safetensors weights."""
+
+    candidates: List[Any] = [
+        getattr(model_meta.model, "model_subtype", None) if model_meta.model else None,
+        getattr(model_meta.model, "name", None) if model_meta.model else None,
+        model_meta.name,
+        model_meta.description,
+        weights_path.stem,
+    ]
+
+    for value in candidates:
+        if not value:
+            continue
+        text = str(value).lower()
+        if "regnet" in text:
+            return "RegNetX800MF"
+        if "efficientnet" in text and "b4" in text:
+            return "EfficientNetB4"
+        if "efficientnet" in text:
+            return "EfficientNetB4"
+
+    logger.warning(
+        "Unable to infer model backbone for %s; defaulting to EfficientNetB4.",
+        weights_path,
+    )
+    return "EfficientNetB4"
+
+
+LEGACY_CLASS_LABELS = [
+    "appendix",
+    "blood",
+    "diverticule",
+    "grasper",
+    "ileocaecalvalve",
+    "ileum",
+    "low_quality",
+    "nbi",
+    "needle",
+    "outside",
+    "polyp",
+    "snare",
+    "water_jet",
+    "wound",
+]
+
+LEGACY_LABEL_ALIASES = {
+    "nbi": "digital_chromo_endoscopy",
+    "grasper": "instrument",
+    "needle": "instrument",
+    "snare": "instrument",
+}
+
+LEGACY_IGNORED_LABELS = {"diverticule"}
+
+
+def _infer_output_classes(weights_path: Path) -> Optional[int]:
+    if weights_path.suffix.lower() != ".safetensors":
+        return None
+
+    try:
+        with safe_open(weights_path, framework="pt", device="cpu") as handle:
+            return int(handle.get_tensor("model.fc.weight").shape[0])
+    except Exception as exc:  # pragma: no cover - defensive logging only
+        logger.debug("Unable to infer output classes from %s: %s", weights_path, exc)
+        return None
+
+
+def _build_label_mapping(source_labels: List[str], target_labels: List[str]) -> Dict[str, List[str]]:
+    if source_labels == target_labels:
+        return {label: [label] for label in target_labels}
+
+    mapping: Dict[str, List[str]] = {label: [] for label in target_labels}
+
+    for label in source_labels:
+        alias = LEGACY_LABEL_ALIASES.get(label, label)
+        if alias in mapping:
+            mapping[alias].append(label)
+        elif label not in LEGACY_IGNORED_LABELS:
+            logger.debug("Label '%s' from source set has no mapping; dropping.", label)
+
+    for label in target_labels:
+        if not mapping[label]:
+            mapping[label] = [label]
+
+    return mapping
+
+
+def _remap_prediction_dict(predictions: Dict[str, Any], mapping: Dict[str, List[str]]) -> Dict[str, Any]:
+    remapped: Dict[str, Any] = {}
+    for target, sources in mapping.items():
+        values: List[Any] = []
+        for source in sources:
+            value = predictions.get(source)
+            if value is not None:
+                values.append(value)
+        if not values:
+            remapped[target] = 0.0
+            continue
+
+        first = values[0]
+        if isinstance(first, np.ndarray):
+            stacked = np.stack(values, axis=0)
+            remapped[target] = stacked.max(axis=0)
+        elif hasattr(first, "__iter__") and not isinstance(first, (float, int)):
+            stacked = np.stack([np.asarray(v) for v in values], axis=0)
+            remapped[target] = stacked.max(axis=0)
+        else:
+            remapped[target] = max(float(v) for v in values)
+
+    return remapped
+
+
+def _extract_text_from_video_frames(video: "VideoFile", frame_fraction: float = 0.001, cap: int = 15) -> Optional[Dict[str, str]]:
     """
     Extracts text from a sample of video frames using OCR based on processor ROIs.
     Requires frames to be extracted. Raises ValueError on pre-condition failure.
@@ -29,7 +181,7 @@ def _extract_text_from_video_frames(
         extract_text_from_rois,
     )  # Local import for dependency isolation
 
-    state = video.get_or_create_state() # Use State helper
+    state = video.get_or_create_state()  # Use State helper
     # --- Pre-condition Check ---
     if not state.frames_extracted:
         # Raise exception
@@ -42,7 +194,7 @@ def _extract_text_from_video_frames(
         raise ValueError(f"Processor not set for video {video.uuid}. Cannot extract text.")
 
     try:
-        frame_paths = video.get_frame_paths() # Use Frame helper
+        frame_paths = video.get_frame_paths()  # Use Frame helper
     except Exception as e:
         logger.error("Error getting frame paths for video %s: %s", video.uuid, e, exc_info=True)
         raise RuntimeError(f"Could not get frame paths for video {video.uuid}") from e
@@ -50,7 +202,7 @@ def _extract_text_from_video_frames(
     n_frames = len(frame_paths)
     if n_frames == 0:
         logger.warning("No frame paths found for video %s during text extraction.", video.uuid)
-        return None # Return None if no frames, not an error condition for this function
+        return None  # Return None if no frames, not an error condition for this function
 
     # Determine number of frames to process
     n_frames_to_process = max(1, int(frame_fraction * n_frames))
@@ -78,10 +230,8 @@ def _extract_text_from_video_frames(
                     rois_texts[roi].append(text)
         except Exception as e:
             # Log error but continue processing other frames
-            logger.error(
-                "Error extracting text from frame %s for video %s: %s", frame_path, video.uuid, e, exc_info=True
-            )
-            errors_encountered = True # Flag that an error occurred
+            logger.error("Error extracting text from frame %s for video %s: %s", frame_path, video.uuid, e, exc_info=True)
+            errors_encountered = True  # Flag that an error occurred
 
     # Determine the most frequent text for each ROI
     most_frequent_texts = {}
@@ -97,17 +247,15 @@ def _extract_text_from_video_frames(
             else:
                 most_frequent_texts[roi] = None
         except Exception as e:
-            logger.error(
-                "Error finding most common text for ROI %s: %s", roi, e, exc_info=True
-            )
+            logger.error("Error finding most common text for ROI %s: %s", roi, e, exc_info=True)
             most_frequent_texts[roi] = None
 
     if errors_encountered:
         logger.warning("Errors occurred during text extraction for some frames of video %s. Results may be incomplete.", video.uuid)
 
     if not most_frequent_texts:
-         logger.info("No text extracted for any ROI for video %s.", video.uuid)
-         return None # Return None if no text found
+        logger.info("No text extracted for any ROI for video %s.", video.uuid)
+        return None  # Return None if no text found
 
     logger.info("Extracted text for video %s: %s", video.uuid, most_frequent_texts)
     return most_frequent_texts
@@ -121,7 +269,7 @@ def _predict_video_pipeline(
     binarize_threshold: float = 0.5,
     test_run: bool = False,
     n_test_frames: int = 10,
-) -> Dict[str, List[Tuple[int, int]]]: # Changed return type to non-optional
+) -> Dict[str, List[Tuple[int, int]]]:  # Changed return type to non-optional
     """
     Executes the video prediction pipeline using an AI model.
     Requires frames to be extracted. Raises exceptions on failure.
@@ -132,13 +280,13 @@ def _predict_video_pipeline(
     """
     # Import heavy dependencies locally
     from ...administration.ai import AiModel
-    try:
-        from ....utils.ai import InferenceDataset, MultiLabelClassificationNet, Classifier
 
+    try:
+        from ....utils.ai import Classifier, InferenceDataset, MultiLabelClassificationNet
         from ....utils.ai.postprocess import (
             concat_pred_dicts,
-            make_smooth_preds,
             find_true_pred_sequences,
+            make_smooth_preds,
         )
     except ImportError as e:
         logger.error("Failed to import endo_ai components: %s. Prediction unavailable.", e, exc_info=True)
@@ -150,7 +298,7 @@ def _predict_video_pipeline(
         n_test_frames = GLOBAL_N_TEST_FRAMES
         logger.info("Using global TEST_RUN settings for prediction pipeline.")
 
-    state = video.get_or_create_state() # Use State helper
+    state = video.get_or_create_state()  # Use State helper
     # --- Pre-condition Check ---
     if not state.frames_extracted:
         # Raise exception
@@ -158,7 +306,7 @@ def _predict_video_pipeline(
     # --- End Pre-condition Check ---
 
     # Frame directory check
-    frame_dir = video.get_frame_dir_path() # Use IO helper
+    frame_dir = video.get_frame_dir_path()  # Use IO helper
     if not frame_dir or not frame_dir.exists() or not any(frame_dir.iterdir()):
         # Raise exception
         raise FileNotFoundError(f"Frame directory {frame_dir} is empty or does not exist for video {video.uuid}. Prediction aborted.")
@@ -180,9 +328,7 @@ def _predict_video_pipeline(
 
     # Get or create VideoPredictionMeta
     try:
-        _video_prediction_meta, created = VideoPredictionMeta.objects.get_or_create(
-            video_file=video, model_meta=model_meta
-        )
+        _video_prediction_meta, created = VideoPredictionMeta.objects.get_or_create(video_file=video, model_meta=model_meta)
         if created:
             logger.info(
                 "Created new VideoPredictionMeta for video %s, model %s.",
@@ -197,12 +343,18 @@ def _predict_video_pipeline(
             )
         # video_prediction_meta.save() # Save is handled by get_or_create
     except Exception as e:
-        logger.error(
-            "Failed to get or create VideoPredictionMeta for video %s, model %s: %s", video.uuid, model_meta.name, e, exc_info=True
-        )
+        logger.error("Failed to get or create VideoPredictionMeta for video %s, model %s: %s", video.uuid, model_meta.name, e, exc_info=True)
         # Raise exception
         raise RuntimeError("Failed to get or create VideoPredictionMeta") from e
 
+    if _is_stub_weights_file(weights_path):
+        logger.info(
+            "Detected stub weights at %s for video %s; skipping model inference and returning empty predictions.",
+            weights_path,
+            video.uuid,
+        )
+        return {}
+
     # --- Dataset Preparation ---
     datasets = {
         "inference_dataset": InferenceDataset,
@@ -214,40 +366,92 @@ def _predict_video_pipeline(
         raise ValueError(f"Dataset class '{dataset_name}' not found for video {video.uuid}. Prediction aborted.")
 
     try:
-        paths = video.get_frame_paths() # Use Frame helper
+        paths = video.get_frame_paths()  # Use Frame helper
         if not paths:
             raise FileNotFoundError(f"No frame paths returned by get_frame_paths for {frame_dir} (Video: {video.uuid})")
     except Exception as e:
-        logger.error(
-            "Error listing or getting frame files from %s for video %s: %s", frame_dir, video.uuid, e, exc_info=True
-        )
+        logger.error("Error listing or getting frame files from %s for video %s: %s", frame_dir, video.uuid, e, exc_info=True)
         raise RuntimeError(f"Error getting frame paths from {frame_dir}") from e
 
     logger.info("Found %d frame files in %s for video %s.", len(paths), frame_dir, video.uuid)
 
-    crop_template = video.get_crop_template() # Use Meta helper
+    crop_template = video.get_crop_template()  # Use Meta helper
     string_paths = [p.as_posix() for p in paths]
-    crops = [crop_template] * len(paths) # Assuming same crop for all frames
+    crops = [crop_template] * len(paths)  # Assuming same crop for all frames
 
     if test_run:
         logger.info("TEST RUN: Using first %d frames for video %s.", n_test_frames, video.uuid)
         string_paths = string_paths[:n_test_frames]
         crops = crops[:n_test_frames]
         if not string_paths:
-             # Raise exception
+            # Raise exception
             raise ValueError(f"Not enough frames ({len(paths)}) for test run (required {n_test_frames}) for video {video.uuid}.")
 
+    label_names = _resolve_label_names(model_meta)
+    if not label_names:
+        raise ValueError(f"Label set '{getattr(model_meta.labelset, 'name', 'unknown')}' has no labels configured.")
+
+    outputs_hint = _infer_output_classes(weights_path)
+
+    network_labels = label_names
+    if outputs_hint and outputs_hint != len(label_names):
+        if outputs_hint == len(LEGACY_CLASS_LABELS):
+            network_labels = LEGACY_CLASS_LABELS
+            logger.info(
+                "Detected legacy multilabel checkpoint with %d classes; using legacy label ordering.",
+                outputs_hint,
+            )
+        else:
+            logger.warning(
+                "Weights %s expect %d outputs while label set '%s' defines %d labels.",
+                weights_path.name,
+                outputs_hint,
+                getattr(model_meta.labelset, "name", "unknown"),
+                len(label_names),
+            )
+
+    label_mapping = _build_label_mapping(network_labels, label_names)
+
+    load_kwargs: Dict[str, Any] = {}
+    if weights_path.suffix.lower() == ".safetensors":
+        load_kwargs.update(
+            {
+                "labels": network_labels,
+                "model_type": _infer_model_type(model_meta, weights_path),
+                "load_imagenet_weights": False,
+                "strict": False,
+            }
+        )
+
+    classifier_config: Optional[Dict[str, Any]] = None
+
     try:
         ds_config = model_meta.get_inference_dataset_config()
         ds = dataset_model_class(string_paths, crops, config=ds_config)
         logger.info("Created dataset '%s' with %d items for video %s.", dataset_name, len(ds), video.uuid)
         if len(ds) > 0:
-            sample = ds[0] # Get a sample for debugging shape
-            logger.debug("Sample shape: %s", sample.shape)
+            sample = ds[0]  # Get a sample for debugging shape
+            logger.debug("Sample shape: %s", getattr(sample, "shape", None))
+
+        try:
+            activation = ModelMeta.get_activation_function(model_meta.activation)
+        except ValueError:
+            logger.warning(
+                "Unsupported activation '%s' for model %s; falling back to sigmoid.",
+                model_meta.activation,
+                model_meta.name,
+            )
+            activation = ModelMeta.get_activation_function("sigmoid")
+
+        classifier_config = {
+            **ds_config,
+            "batchsize": model_meta.batchsize or 16,
+            "num_workers": model_meta.num_workers or 0,
+            "activation": activation,
+            "labels": network_labels,
+        }
     except Exception as e:
-        logger.error(
-            "Failed to create dataset '%s' for video %s: %s", dataset_name, video.uuid, e, exc_info=True
-        )
+        logger.error("Failed to create dataset '%s' for video %s: %s", dataset_name, video.uuid, e, exc_info=True)
         # Raise exception
         raise RuntimeError(f"Failed to create dataset '{dataset_name}'") from e
 
@@ -255,38 +459,47 @@ def _predict_video_pipeline(
     try:
         # Check if CUDA is available
         import torch
+
         if torch.cuda.is_available():
             try:
-                # Try loading on GPU first
+                device = torch.device("cuda")
                 ai_model_instance = MultiLabelClassificationNet.load_from_checkpoint(
-                    checkpoint_path=weights_path.as_posix(), # Ensure path is string
+                    checkpoint_path=weights_path.as_posix(),
+                    map_location=device,
+                    **load_kwargs,
                 )
-                # Attempt to move to GPU
-                _ = ai_model_instance.cuda()
+                ai_model_instance = ai_model_instance.to(device)
                 logger.info("Loaded model on GPU for video %s.", video.uuid)
             except RuntimeError as cuda_err:
-                # If GPU loading fails, fall back to CPU
-                logger.warning("GPU loading failed for video %s: %s. Falling back to CPU.", video.uuid, cuda_err)
+                logger.warning(
+                    "GPU loading failed for video %s: %s. Falling back to CPU.",
+                    video.uuid,
+                    cuda_err,
+                )
+                device = torch.device("cpu")
                 ai_model_instance = MultiLabelClassificationNet.load_from_checkpoint(
                     checkpoint_path=weights_path.as_posix(),
-                    map_location='cpu'
+                    map_location=device,
+                    **load_kwargs,
                 )
+                ai_model_instance = ai_model_instance.to(device)
                 logger.info("Loaded model on CPU for video %s.", video.uuid)
         else:
             # No CUDA available, load directly on CPU
             logger.info("CUDA not available. Loading model on CPU for video %s.", video.uuid)
+            device = torch.device("cpu")
             ai_model_instance = MultiLabelClassificationNet.load_from_checkpoint(
                 checkpoint_path=weights_path.as_posix(),
-                map_location='cpu'
+                map_location=device,
+                **load_kwargs,
             )
+            ai_model_instance = ai_model_instance.to(device)
 
-        _ = ai_model_instance.eval() # Set to evaluation mode
-        classifier = Classifier(ai_model_instance, verbose=True) # Assuming Classifier exists
+        _ = ai_model_instance.eval()  # Set to evaluation mode
+        classifier = Classifier(ai_model_instance, config=classifier_config or {}, verbose=True)
         logger.info("AI model loaded successfully for video %s from %s.", video.uuid, weights_path)
     except Exception as e:
-        logger.error(
-            "Failed to load AI model for video %s from %s: %s", video.uuid, weights_path, e, exc_info=True
-        )
+        logger.error("Failed to load AI model for video %s from %s: %s", video.uuid, weights_path, e, exc_info=True)
         # Raise exception
         raise RuntimeError(f"Failed to load AI model from {weights_path}") from e
 
@@ -299,15 +512,19 @@ def _predict_video_pipeline(
         logger.error("Inference failed for video %s: %s", video.uuid, e, exc_info=True)
         # CUDA-OOM Fallback: Speicher freigeben und CPU versuchen
         try:
-            import torch, gc
-            is_oom = isinstance(e, (getattr(torch.cuda, 'OutOfMemoryError', RuntimeError), RuntimeError)) and (
-                'out of memory' in str(e).lower() or 'cuda out of memory' in str(e).lower()
+            import gc
+
+            import torch
+
+            is_oom = isinstance(e, (getattr(torch.cuda, "OutOfMemoryError", RuntimeError), RuntimeError)) and (
+                "out of memory" in str(e).lower() or "cuda out of memory" in str(e).lower()
             )
         except Exception:
             is_oom = False
-        if 'torch' in globals() or 'torch' in locals():
+        if "torch" in globals() or "torch" in locals():
             try:
                 import torch  # ensure available in this scope
+
                 if torch.cuda.is_available() and is_oom:
                     logger.warning("CUDA OOM detected. Freeing CUDA cache and retrying on CPU…")
                     try:
@@ -339,17 +556,20 @@ def _predict_video_pipeline(
     try:
         logger.info("Post-processing predictions for video %s...", video.uuid)
         readable_predictions = [classifier.readable(p) for p in predictions]
+        if label_mapping:
+            readable_predictions = [_remap_prediction_dict(prediction, label_mapping) for prediction in readable_predictions]
 
         merged_predictions = concat_pred_dicts(readable_predictions)
 
-        fps = video.get_fps() # Use Meta helper
+        fps = video.get_fps()  # Use Meta helper
         if not fps:
             logger.warning(
                 "Video FPS is unknown for %s. Smoothing/sequence calculations might be inaccurate. Using default 30 FPS.",
                 video.uuid,
             )
-            fps = 30 # Default FPS if unknown
+            fps = 30  # Default FPS if unknown
 
+        fps = int(fps)
         smooth_merged_predictions = {}
         for key in merged_predictions.keys():
             smooth_merged_predictions[key] = make_smooth_preds(
@@ -360,9 +580,7 @@ def _predict_video_pipeline(
 
         binary_smooth_merged_predictions = {}
         for key in smooth_merged_predictions.keys():
-            binary_smooth_merged_predictions[key] = (
-                smooth_merged_predictions[key] > binarize_threshold
-            )
+            binary_smooth_merged_predictions[key] = smooth_merged_predictions[key] > binarize_threshold
 
         sequences = {}
         for label, prediction_array in binary_smooth_merged_predictions.items():
@@ -390,10 +608,10 @@ def _predict_video_entry(
     binarize_threshold: float = 0.5,
     test_run: bool = GLOBAL_TEST_RUN,
     n_test_frames: int = GLOBAL_N_TEST_FRAMES,
-    save_results: bool = True, # Note: save_results is handled in video_file.py now
+    save_results: bool = True,  # Note: save_results is handled in video_file.py now
 ):
     """Entry point called from VideoFile.predict_video. Imports and calls the main prediction logic."""
-    from endoreg_db.models import AiModel, ModelMeta # Local import
+    from endoreg_db.models import AiModel, ModelMeta  # Local import
 
     try:
         ai_model = AiModel.objects.get(name=model_name)
@@ -412,7 +630,7 @@ def _predict_video_entry(
     # --- Explicitly pass only the arguments expected by _predict_video_pipeline ---
     predicted_sequences = _predict_video_pipeline(
         video=video,
-        model_meta=model_meta, # Pass the fetched ModelMeta object
+        model_meta=model_meta,  # Pass the fetched ModelMeta object
         dataset_name=dataset_name,
         smooth_window_size_s=smooth_window_size_s,
         binarize_threshold=binarize_threshold,
@@ -425,15 +643,11 @@ def _predict_video_entry(
     return predicted_sequences, model_meta
 
 
-def _extract_text_information(
-    video: "VideoFile", frame_fraction: float = 0.001, cap: int = 15
-) -> Optional[Dict[str, str]]:
+def _extract_text_information(video: "VideoFile", frame_fraction: float = 0.001, cap: int = 15) -> Optional[Dict[str, str]]:
     """Facade function to call the text extraction logic."""
     logger.info("Attempting text extraction for video %s.", video.uuid)
 
-    extracted_data = _extract_text_from_video_frames(
-        video=video, frame_fraction=frame_fraction, cap=cap
-    )
+    extracted_data = _extract_text_from_video_frames(video=video, frame_fraction=frame_fraction, cap=cap)
 
     if extracted_data is not None:
         logger.info("Text extraction successful for video %s.", video.uuid)