dragon-ml-toolbox 13.3.0__py3-none-any.whl → 16.2.0__py3-none-any.whl

ml_tools/ML_vision_evaluation.py

@@ -0,0 +1,284 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+import torch
+from sklearn.metrics import (
+    accuracy_score,
+    f1_score,
+    jaccard_score,
+    confusion_matrix,
+    ConfusionMatrixDisplay
+)
+from pathlib import Path
+from typing import Union, Optional, List, Dict
+import json
+from torchmetrics.detection import MeanAveragePrecision
+
+from .path_manager import make_fullpath
+from ._logger import _LOGGER
+from ._script_info import _script_info
+from ._keys import VisionKeys
+from .ML_configuration import (BinarySegmentationMetricsFormat,
+                               MultiClassSegmentationMetricsFormat,
+                               _BaseSegmentationFormat)
+
+
+__all__ = [
+    "segmentation_metrics",
+    "object_detection_metrics"
+]
+
+DPI_value = 250
+
+
+def segmentation_metrics(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    save_dir: Union[str, Path],
+    class_names: Optional[List[str]] = None,
+    config: Optional[Union[BinarySegmentationMetricsFormat, MultiClassSegmentationMetricsFormat]] = None
+):
+    """
+    Calculates and saves pixel-level metrics for segmentation tasks.
+
+    Metrics include Pixel Accuracy, Dice (F1-score), and IoU (Jaccard).
+    It calculates 'micro', 'macro', and 'weighted' averages and saves
+    a pixel-level confusion matrix and a metrics heatmap.
+    
+    Note: This function expects integer-based masks (e.g., shape [N, H, W] or [H, W]),
+    not one-hot encoded masks.
+    
+    Args:
+        y_true (np.ndarray): Ground truth masks (e.g., shape [N, H, W]).
+        y_pred (np.ndarray): Predicted masks (e.g., shape [N, H, W]).
+        save_dir (str | Path): Directory to save the metrics report and plots.
+        class_names (List[str] | None): Names of the classes for the report.
+        config (object): Formatting configuration object.
+    """
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+    
+    # --- Parse Config or use defaults ---
+    if config is None:
+        format_config = _BaseSegmentationFormat()
+    else:
+        format_config = config
+
+    # --- Set Matplotlib font size ---
+    original_rc_params = plt.rcParams.copy()
+    plt.rcParams.update({'font.size': format_config.font_size})
+    
+    # Get all unique class labels present in either true or pred
+    labels = np.unique(np.concatenate((np.unique(y_true), np.unique(y_pred)))).astype(int)
+    
+    # --- Setup Class Names ---
+    display_names = []
+    if class_names is None:
+        display_names = [f"Class {i}" for i in labels]
+    else:
+        if len(class_names) != len(labels):
+            _LOGGER.warning(f"Number of class_names ({len(class_names)}) does not match number of unique labels ({len(labels)}). Using default names.")
+            display_names = [f"Class {i}" for i in labels]
+        else:
+            display_names = class_names
+
+    # Flatten masks for sklearn metrics
+    y_true_flat = y_true.ravel()
+    y_pred_flat = y_pred.ravel()
+
+    _LOGGER.info("--- Calculating Segmentation Metrics ---")
+
+    # --- 1. Calculate Metrics ---
+    pix_acc = accuracy_score(y_true_flat, y_pred_flat)
+    
+    # Calculate all average types
+    dice_micro = f1_score(y_true_flat, y_pred_flat, average='micro', labels=labels)
+    iou_micro = jaccard_score(y_true_flat, y_pred_flat, average='micro', labels=labels)
+    
+    dice_macro = f1_score(y_true_flat, y_pred_flat, average='macro', labels=labels, zero_division=0)
+    iou_macro = jaccard_score(y_true_flat, y_pred_flat, average='macro', labels=labels, zero_division=0)
+    
+    dice_weighted = f1_score(y_true_flat, y_pred_flat, average='weighted', labels=labels, zero_division=0)
+    iou_weighted = jaccard_score(y_true_flat, y_pred_flat, average='weighted', labels=labels, zero_division=0)
+    
+    # Per-class metrics
+    dice_per_class = f1_score(y_true_flat, y_pred_flat, average=None, labels=labels, zero_division=0)
+    iou_per_class = jaccard_score(y_true_flat, y_pred_flat, average=None, labels=labels, zero_division=0)
+    
+    # --- 2. Create and Save Report ---
+    report_lines = [
+        "--- Segmentation Report ---",
+        f"\nOverall Pixel Accuracy: {pix_acc:.4f}\n",
+        "--- Averaged Metrics ---",
+        f"{'Average':<10} | {'Dice (F1)':<12} | {'IoU (Jaccard)':<12}",
+        "-"*41,
+        f"{'Micro':<10} | {dice_micro:<12.4f} | {iou_micro:<12.4f}",
+        f"{'Macro':<10} | {dice_macro:<12.4f} | {iou_macro:<12.4f}",
+        f"{'Weighted':<10} | {dice_weighted:<12.4f} | {iou_weighted:<12.4f}",
+        "\n--- Per-Class Metrics ---",
+    ]
+
+    per_class_data = {
+        'Class': display_names,
+        'Dice': dice_per_class,
+        'IoU': iou_per_class
+    }
+    per_class_df = pd.DataFrame(per_class_data)
+    report_lines.append(per_class_df.to_string(index=False, float_format="%.4f"))
+
+    report_string = "\n".join(report_lines)
+    # print(report_string) # <-- I removed the print(report_string)
+    
+    # Save text report
+    save_filename = VisionKeys.SEGMENTATION_REPORT + ".txt"
+    report_path = save_dir_path / save_filename
+    report_path.write_text(report_string, encoding="utf-8")
+    _LOGGER.info(f"📝 Segmentation report saved as '{report_path.name}'")
+
+    # --- 3. Save Per-Class Metrics Heatmap ---
+    try:
+        plt.figure(figsize=(max(8, len(labels) * 0.5), 6), dpi=DPI_value)
+        sns.heatmap(
+            per_class_df.set_index('Class').T, 
+            annot=True, 
+            cmap=format_config.heatmap_cmap, # Use config cmap
+            fmt='.3f',
+            linewidths=0.5
+        )
+        plt.title("Per-Class Segmentation Metrics")
+        plt.tight_layout()
+        heatmap_filename = VisionKeys.SEGMENTATION_HEATMAP + ".svg"
+        heatmap_path = save_dir_path / heatmap_filename
+        plt.savefig(heatmap_path)
+        _LOGGER.info(f"📊 Metrics heatmap saved as '{heatmap_path.name}'")
+        plt.close()
+    except Exception as e:
+        _LOGGER.error(f"Could not generate segmentation metrics heatmap: {e}")
+
+    # --- 4. Save Pixel-level Confusion Matrix ---
+    try:
+        # Calculate CM
+        cm = confusion_matrix(y_true_flat, y_pred_flat, labels=labels)
+        
+        # Plot
+        fig_cm, ax_cm = plt.subplots(figsize=(max(8, len(labels) * 0.8), max(8, len(labels) * 0.8)), dpi=100)
+        disp = ConfusionMatrixDisplay(
+            confusion_matrix=cm,
+            display_labels=display_names
+        )
+        disp.plot(cmap=format_config.cm_cmap, ax=ax_cm, xticks_rotation=45) # Use config cmap
+        
+        # Manually update font size of cell texts
+        for text in disp.text_.flatten(): # type: ignore
+            text.set_fontsize(format_config.font_size)
+        
+        ax_cm.set_title("Pixel-Level Confusion Matrix")
+        plt.tight_layout()
+        segmentation_cm_filename = VisionKeys.SEGMENTATION_CONFUSION_MATRIX + ".svg"
+        cm_path = save_dir_path / segmentation_cm_filename
+        plt.savefig(cm_path)
+        _LOGGER.info(f"❇️ Pixel-level confusion matrix saved as '{cm_path.name}'")
+        plt.close(fig_cm)
+    except Exception as e:
+        _LOGGER.error(f"Could not generate confusion matrix: {e}")
+        
+    # --- Restore RC params ---
+    plt.rcParams.update(original_rc_params)
+
+
+def object_detection_metrics(
+    preds: List[Dict[str, torch.Tensor]],
+    targets: List[Dict[str, torch.Tensor]],
+    save_dir: Union[str, Path],
+    class_names: Optional[List[str]] = None,
+    print_output: bool=False
+):
+    """
+    Calculates and saves object detection metrics (mAP) using torchmetrics.
+
+    This function expects predictions and targets in the standard
+    torchvision format (list of dictionaries).
+
+    Args:
+        preds (List[Dict[str, torch.Tensor]]): A list of predictions.
+            Each dict must contain:
+            - 'boxes': [N, 4] (xmin, ymin, xmax, ymax)
+            - 'scores': [N]
+            - 'labels': [N]
+        targets (List[Dict[str, torch.Tensor]]): A list of ground truths.
+            Each dict must contain:
+            - 'boxes': [M, 4]
+            - 'labels': [M]
+        save_dir (str | Path): Directory to save the metrics report (as JSON).
+        class_names (List[str] | None): A list of class names, including 'background'
+            at index 0. Used to label per-class metrics in the report.
+        print_output (bool): If True, prints the JSON report to the console.
+    """
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+
+    _LOGGER.info("--- Calculating Object Detection Metrics (mAP) ---")
+
+    try:
+        # Initialize the metric with standard COCO settings
+        metric = MeanAveragePrecision(box_format='xyxy')
+        
+        # Move preds and targets to the same device (e.g., CPU for metric calculation)
+        # This avoids device mismatches if model was on GPU
+        device = torch.device("cpu")
+        preds_cpu = [{k: v.to(device) for k, v in p.items()} for p in preds]
+        targets_cpu = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        
+        # Update the metric
+        metric.update(preds_cpu, targets_cpu)
+        
+        # Compute the final metrics
+        results = metric.compute()
+        
+        # --- Handle class names for per-class metrics ---
+        report_class_names = None
+        if class_names:
+            if class_names[0].lower() in ['background', "bg"]:
+                report_class_names = class_names[1:] # Skip background (class 0)
+            else:
+                _LOGGER.warning("class_names provided to object_detection_metrics, but 'background' was not class 0. Using all provided names.")
+                report_class_names = class_names
+        
+        # Convert all torch tensors in results to floats/lists for JSON serialization
+        serializable_results = {}
+        for key, value in results.items():
+            if isinstance(value, torch.Tensor):
+                if value.numel() == 1:
+                    serializable_results[key] = value.item()
+                # Check if it's a 1D tensor, we have class names, and it's a known per-class key
+                elif value.ndim == 1 and report_class_names and key in ('map_per_class', 'mar_100_per_class', 'mar_1_per_class', 'mar_10_per_class'):
+                    per_class_list = value.cpu().numpy().tolist()
+                    # Map names to values
+                    if len(per_class_list) == len(report_class_names):
+                        serializable_results[key] = {name: val for name, val in zip(report_class_names, per_class_list)}
+                    else:
+                        _LOGGER.warning(f"Length mismatch for '{key}': {len(per_class_list)} values vs {len(report_class_names)} class names. Saving as raw list.")
+                        serializable_results[key] = per_class_list
+                else:
+                    serializable_results[key] = value.cpu().numpy().tolist()
+            else:
+                serializable_results[key] = value
+        
+        # Pretty print to console
+        if print_output:
+            print(json.dumps(serializable_results, indent=4))
+
+        # Save JSON report
+        detection_report_filename = VisionKeys.OBJECT_DETECTION_REPORT + ".json"
+        report_path = save_dir_path / detection_report_filename
+        with open(report_path, 'w') as f:
+            json.dump(serializable_results, f, indent=4)
+        
+        _LOGGER.info(f"📊 Object detection (mAP) report saved as '{report_path.name}'")
+
+    except Exception as e:
+        _LOGGER.error(f"Failed to compute mAP: {e}")
+        raise
+
+
+def info():
+    _script_info(__all__)

ml_tools/ML_vision_inference.py

@@ -0,0 +1,405 @@
+import torch
+from torch import nn
+import numpy as np  #numpy array return value
+from pathlib import Path
+from typing import Union, Literal, Dict, Any, List, Optional, Callable
+from PIL import Image
+from torchvision import transforms
+
+from ._script_info import _script_info
+from ._logger import _LOGGER
+from ._keys import PyTorchInferenceKeys, MLTaskKeys
+from .ML_vision_transformers import _load_recipe_and_build_transform
+from .ML_inference import _BaseInferenceHandler
+
+
+__all__ = [
+    "DragonVisionInferenceHandler"
+]
+
+
+class DragonVisionInferenceHandler(_BaseInferenceHandler):
+    """
+    Handles loading a PyTorch vision model's state dictionary and performing inference.
+
+    This class is specifically for vision models, which typically expect
+    4D Tensors (B, C, H, W) or Lists of Tensors as input.
+    """
+    def __init__(self,
+                 model: nn.Module,
+                 state_dict: Union[str, Path],
+                 task: Literal["binary image classification", "multiclass image classification", "binary segmentation", "multiclass segmentation", "object detection"],
+                 device: str = 'cpu',
+                 transform_source: Optional[Union[str, Path, Callable]] = None):
+        """
+        Initializes the vision inference handler.
+
+        Args:
+            model (nn.Module): An instantiated PyTorch model from ML_vision_models.
+            state_dict (str | Path): Path to the saved .pth model state_dict file.
+            task (str): The type of vision task.
+            device (str): The device to run inference on ('cpu', 'cuda', 'mps').
+            transform_source (str | Path | Callable | None): 
+                - A path to a .json recipe file (str or Path).
+                - A pre-built transformation pipeline (Callable).
+                - None, in which case .set_transform() must be called explicitly to set transformations.
+        
+        Note: class_map (Dict[int, str]) will be loaded from the model file, to set or override it use `.set_class_map()`.
+        """
+        super().__init__(model, state_dict, device, None)
+
+        self._transform: Optional[Callable] = None
+        self._is_transformed: bool = False
+
+        if task not in [MLTaskKeys.BINARY_IMAGE_CLASSIFICATION, MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION, MLTaskKeys.BINARY_SEGMENTATION, MLTaskKeys.MULTICLASS_SEGMENTATION, MLTaskKeys.OBJECT_DETECTION]:
+            _LOGGER.error(f"Unsupported task: '{task}'.")
+            raise ValueError()
+        self.task = task
+        
+        self.expected_in_channels: int = 3 # Default to RGB
+        if hasattr(model, 'in_channels'):
+            self.expected_in_channels = model.in_channels # type: ignore
+            _LOGGER.info(f"Model expects {self.expected_in_channels} input channels.")
+        else:
+            _LOGGER.warning("Could not determine 'in_channels' from model. Defaulting to 3 (RGB). Modify with '.expected_in_channels'.")
+        
+        if transform_source:
+            self.set_transform(transform_source)
+            self._is_transformed = True
+
+    def _preprocess_batch(self, inputs: Union[torch.Tensor, List[torch.Tensor]]) -> Union[torch.Tensor, List[torch.Tensor]]:
+        """
+        Validates input and moves it to the correct device.
+        - For Classification/Segmentation: Expects 4D Tensor (B, C, H, W).
+        - For Object Detection: Expects List[Tensor(C, H, W)].
+        """
+        if self.task == MLTaskKeys.OBJECT_DETECTION:
+            if not isinstance(inputs, list) or not all(isinstance(t, torch.Tensor) for t in inputs):
+                _LOGGER.error("Input for object_detection must be a List[torch.Tensor].")
+                raise ValueError("Invalid input type for object detection.")
+            # Move each tensor in the list to the device
+            return [t.float().to(self.device) for t in inputs]
+        
+        else: # Classification or Segmentation
+            if not isinstance(inputs, torch.Tensor):
+                _LOGGER.error(f"Input for {self.task} must be a torch.Tensor.")
+                raise ValueError(f"Invalid input type for {self.task}.")
+                
+            if inputs.ndim != 4: # type: ignore
+                 _LOGGER.error(f"Input tensor for {self.task} must be 4D (B, C, H, W). Got {inputs.ndim}D.") # type: ignore
+                 raise ValueError("Input tensor must be 4D.")
+            
+            return inputs.float().to(self.device) 
+        
+    def set_transform(self, transform_source: Union[str, Path, Callable]):
+        """
+        Sets or updates the inference transformation pipeline from a recipe file or a direct Callable.
+
+        Args:
+            transform_source (str, Path, Callable):
+                - A path to a .json recipe file (str or Path).
+                - A pre-built transformation pipeline (Callable).
+        """
+        if self._is_transformed:
+            _LOGGER.warning("Transformations were previously applied. Applying new transformations...")
+            
+        if isinstance(transform_source, (str, Path)):
+            _LOGGER.info(f"Loading transform from recipe file: '{transform_source}'")
+            try:
+                # Use the loader function
+                self._transform = _load_recipe_and_build_transform(transform_source)
+            except Exception as e:
+                _LOGGER.error(f"Failed to load transform from recipe '{transform_source}': {e}")
+                raise
+        elif isinstance(transform_source, Callable):
+            _LOGGER.info("Inference transform has been set from a direct Callable.")
+            self._transform = transform_source
+        else:
+            _LOGGER.error(f"Invalid transform_source type: {type(transform_source)}. Must be str, Path, or Callable.")
+            raise TypeError("transform_source must be a file path or a Callable.")
+
+    def predict_batch(self, inputs: Union[torch.Tensor, List[torch.Tensor]]) -> Dict[str, Any]:
+        """
+        Core batch prediction method for vision models.
+        All preprocessing (resizing, normalization) should be done *before* calling this method.
+
+        Args:
+            inputs (torch.Tensor | List[torch.Tensor]):
+                - For binary/multiclass image classification or binary/multiclass image segmentation tasks, 
+                  a 4D torch.Tensor (B, C, H, W).
+                - For 'object_detection', a List of 3D torch.Tensors 
+                  [(C, H, W), ...], each with its own size.
+
+        Returns:
+            A dictionary containing the output tensors.
+            - Classification: {labels, probabilities}
+            - Segmentation: {labels, probabilities} (labels is the mask)
+            - Object Detection: {predictions} (List of dicts)
+        """
+        
+        processed_inputs = self._preprocess_batch(inputs)
+        
+        with torch.no_grad():
+            # get outputs
+            output = self.model(processed_inputs)
+            if self.task == MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION:
+                # process
+                probs = torch.softmax(output, dim=1)
+                labels = torch.argmax(probs, dim=1)
+                return {
+                    PyTorchInferenceKeys.LABELS: labels,       # (B,)
+                    PyTorchInferenceKeys.PROBABILITIES: probs  # (B, num_classes)
+                }
+            
+            elif self.task == MLTaskKeys.BINARY_IMAGE_CLASSIFICATION:
+                # Assumes model output is [N, 1] (a single logit)
+                # Squeeze output from [N, 1] to [N] if necessary
+                if output.ndim == 2 and output.shape[1] == 1:
+                    output = output.squeeze(1)
+                    
+                probs = torch.sigmoid(output) # Probability of positive class
+                labels = (probs >= self._classification_threshold).int()
+                return {
+                    PyTorchInferenceKeys.LABELS: labels,       
+                    PyTorchInferenceKeys.PROBABILITIES: probs  
+                }
+            
+            elif self.task == MLTaskKeys.BINARY_SEGMENTATION:
+                # Assumes model output is [N, 1, H, W] (logits for positive class)
+                probs = torch.sigmoid(output) # Shape [N, 1, H, W]
+                labels = (probs >= self._classification_threshold).int() # Shape [N, 1, H, W]
+                return {
+                    PyTorchInferenceKeys.LABELS: labels,       
+                    PyTorchInferenceKeys.PROBABILITIES: probs  
+                }
+                
+            elif self.task == MLTaskKeys.MULTICLASS_SEGMENTATION:
+                # output shape [N, C, H, W]
+                probs = torch.softmax(output, dim=1)
+                labels = torch.argmax(probs, dim=1) # shape [N, H, W]
+                return {
+                    PyTorchInferenceKeys.LABELS: labels,       # (N, H, W)
+                    PyTorchInferenceKeys.PROBABILITIES: probs  # (N, num_classes, H, W)
+                }
+
+            elif self.task == MLTaskKeys.OBJECT_DETECTION:
+                return {
+                    PyTorchInferenceKeys.PREDICTIONS: output
+                }
+            
+            else:
+                # This should be unreachable due to validation
+                raise ValueError(f"Unknown task: {self.task}")
+
+    def predict(self, single_input: torch.Tensor) -> Dict[str, Any]:
+        """
+        Core single-sample prediction method for vision models.
+        All preprocessing (resizing, normalization) should be done *before*
+        calling this method.
+
+        Args:
+            single_input (torch.Tensor):
+                - A 3D torch.Tensor (C, H, W) for any task.
+
+        Returns:
+            A dictionary containing the output tensors for a single sample.
+            - Classification: {labels, probabilities} (label is 0-dim)
+            - Segmentation: {labels, probabilities} (label is a 2D (multiclass) or 3D (binary) mask)
+            - Object Detection: {boxes, labels, scores} (single dict)
+        """
+        if not isinstance(single_input, torch.Tensor) or single_input.ndim != 3:
+             _LOGGER.error(f"Input for predict() must be a 3D tensor (C, H, W). Got {single_input.ndim}D.")
+             raise ValueError()
+        
+        # --- 1. Batch the input based on task ---
+        if self.task == MLTaskKeys.OBJECT_DETECTION:
+            batched_input = [single_input] # List of one tensor
+        else:
+            batched_input = single_input.unsqueeze(0)
+
+        # --- 2. Call batch prediction ---
+        batch_results = self.predict_batch(batched_input)
+
+        # --- 3. Un-batch the results ---
+        if self.task == MLTaskKeys.OBJECT_DETECTION:
+            # batch_results['predictions'] is a List[Dict]. We want the first (and only) Dict.
+            return batch_results[PyTorchInferenceKeys.PREDICTIONS][0]
+        else:
+            # 'labels' and 'probabilities' are tensors. Get the 0-th element.
+            # (B, ...) -> (...)
+            single_results = {key: value[0] for key, value in batch_results.items()}
+            return single_results
+
+    # --- NumPy Convenience Wrappers (on CPU) ---
+
+    def predict_batch_numpy(self, inputs: Union[torch.Tensor, List[torch.Tensor]]) -> Dict[str, Any]:
+        """
+        Convenience wrapper for predict_batch that returns NumPy arrays. With Labels if set.
+        
+        Returns:
+            Dict: A dictionary containing the outputs as NumPy arrays.
+            - Obj. Detection: {predictions: List[Dict[str, np.ndarray]]}
+            - Classification: {labels: np.ndarray, label_names: List[str], probabilities: np.ndarray}
+            - Segmentation: {labels: np.ndarray, probabilities: np.ndarray}
+        """
+        tensor_results = self.predict_batch(inputs)
+        
+        if self.task == MLTaskKeys.OBJECT_DETECTION:
+            # Output is List[Dict[str, Tensor]]
+            # Convert each tensor inside each dict to numpy
+            numpy_results = []
+            for pred_dict in tensor_results[PyTorchInferenceKeys.PREDICTIONS]:
+                # Convert all tensors to numpy
+                np_dict = {key: value.cpu().numpy() for key, value in pred_dict.items()}
+                
+                # 3D pixel to string map unnecessary 
+                # if self._idx_to_class and PyTorchInferenceKeys.LABELS in np_dict:
+                #     np_dict[PyTorchInferenceKeys.LABEL_NAMES] = [
+                #         self._idx_to_class.get(label_id, "Unknown") 
+                #         for label_id in np_dict[PyTorchInferenceKeys.LABELS]
+                #     ]
+                numpy_results.append(np_dict)
+            return {PyTorchInferenceKeys.PREDICTIONS: numpy_results}
+        
+        else:
+            # Output is Dict[str, Tensor] (for Classification or Segmentation)
+            numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
+            
+            # Add string names for classification if map exists
+            is_image_classification = self.task in [
+                MLTaskKeys.BINARY_IMAGE_CLASSIFICATION, 
+                MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION
+            ]
+            
+            if is_image_classification and self._idx_to_class and PyTorchInferenceKeys.LABELS in numpy_results:
+                int_labels = numpy_results[PyTorchInferenceKeys.LABELS] # This is a (B,) array
+                numpy_results[PyTorchInferenceKeys.LABEL_NAMES] = [
+                    self._idx_to_class.get(label_id, "Unknown")
+                    for label_id in int_labels
+                ]
+            
+            return numpy_results
+
+    def predict_numpy(self, single_input: torch.Tensor) -> Dict[str, Any]:
+        """
+        Convenience wrapper for predict that returns NumPy arrays/scalars.
+
+        Returns:
+            Dict: A dictionary containing the outputs as NumPy arrays/scalars.
+            - Obj. Detection: {boxes: np.ndarray, labels: np.ndarray, scores: np.ndarray, label_names: List[str]}
+            - Classification: {labels: int, label_names: str, probabilities: np.ndarray}
+            - Segmentation: {labels: np.ndarray, probabilities: np.ndarray}
+        """
+        tensor_results = self.predict(single_input)
+        
+        if self.task == MLTaskKeys.OBJECT_DETECTION:
+            # Output is Dict[str, Tensor]
+            # Convert each tensor to numpy
+            numpy_results = {
+                key: value.cpu().numpy() for key, value in tensor_results.items()
+            }
+            
+            # Add string names if map exists
+            # if self._idx_to_class and PyTorchInferenceKeys.LABELS in numpy_results:
+            #     int_labels = numpy_results[PyTorchInferenceKeys.LABELS]
+                
+            #     numpy_results[PyTorchInferenceKeys.LABEL_NAMES] = [
+            #         self._idx_to_class.get(label_id, "Unknown")
+            #         for label_id in int_labels
+            #     ]
+                
+            return numpy_results
+            
+        elif self.task in [MLTaskKeys.BINARY_IMAGE_CLASSIFICATION, MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION]:
+            # Output is Dict[str, Tensor(0-dim) or Tensor(1-dim)]
+            int_label = tensor_results[PyTorchInferenceKeys.LABELS].item()
+            label_name = "Unknown"
+            if self._idx_to_class:
+                label_name = self._idx_to_class.get(int_label, "Unknown")
+
+            return {
+                PyTorchInferenceKeys.LABELS: int_label,
+                PyTorchInferenceKeys.LABEL_NAMES: label_name,
+                PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
+            }
+        else: # image_segmentation (binary or multiclass)
+            # Output is Dict[str, Tensor(2D) or Tensor(3D)]
+            return {
+                PyTorchInferenceKeys.LABELS: tensor_results[PyTorchInferenceKeys.LABELS].cpu().numpy(),
+                PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
+            }
+            
+    def predict_from_pil(self, image: Image.Image) -> Dict[str, Any]:
+        """
+        Applies the stored transform to a single PIL image and returns the prediction.
+
+        Args:
+            image (PIL.Image.Image): The input PIL image.
+
+        Returns:
+            Dict: A dictionary containing the prediction results. See `predict_numpy()` for task-specific output structures.
+        """
+        if self._transform is None:
+            _LOGGER.error("Cannot predict from PIL image: No transform has been set. Call .set_transform() or provide transform_source in __init__.")
+            raise RuntimeError("Inference transform is not set.")
+
+        # Apply the transformation pipeline (e.g., resize, crop, ToTensor, normalize)
+        try:
+            transformed_image = self._transform(image)
+        except Exception as e:
+            _LOGGER.error(f"Error applying transform to PIL image: {e}")
+            raise
+            
+        # --- Validation ---
+        if not isinstance(transformed_image, torch.Tensor):
+            _LOGGER.error("The provided transform did not return a torch.Tensor. Does it include transforms.ToTensor()?")
+            raise ValueError("Transform pipeline must output a torch.Tensor.")
+            
+        if transformed_image.ndim != 3:
+            _LOGGER.warning(f"Expected transform to output a 3D (C, H, W) tensor, but got {transformed_image.ndim}D. Attempting to proceed.")
+            # .predict_numpy() -> .predict() which expects a 3D tensor
+            if transformed_image.ndim == 4 and transformed_image.shape[0] == 1:
+                transformed_image = transformed_image.squeeze(0) # Fix if user's transform adds a batch dim
+                _LOGGER.warning("Removed an extra batch dimension.")
+            else:
+                raise ValueError(f"Transform must output a 3D (C, H, W) tensor, got {transformed_image.shape}.")
+
+        # Use the existing single-item predict method
+        return self.predict_numpy(transformed_image)
+
+    def predict_from_file(self, image_path: Union[str, Path]) -> Dict[str, Any]:
+        """
+        Loads a single image from a file, applies the stored transform, and returns the prediction.
+
+        This is a convenience wrapper that loads the image and calls `predict_from_pil()`.
+
+        Args:
+            image_path (str | Path): The file path to the input image.
+
+        Returns:
+            Dict: A dictionary containing the prediction results. See `predict_numpy()` for task-specific output structures.
+        """
+        try:
+            # --- Use expected_in_channels to set PIL mode ---
+            pil_mode: str
+            if self.expected_in_channels == 1:
+                pil_mode = "L"  # Grayscale
+            elif self.expected_in_channels == 4:
+                pil_mode = "RGBA" # RGB + Alpha
+            else:
+                if self.expected_in_channels != 3: # 2, 5+ channels not supported by PIL convert
+                    _LOGGER.warning(f"Model expects {self.expected_in_channels} channels. PIL conversion is limited, defaulting to 3 channels (RGB). The transformations must convert it to the desired channel dimensions.")
+                # Default to RGB. If 2-channels are needed, the transform recipe *must* be responsible for handling the conversion from a 3-channel PIL image.
+                pil_mode = "RGB"
+                
+            image = Image.open(image_path).convert(pil_mode)
+        except Exception as e:
+            _LOGGER.error(f"Failed to load and convert image from '{image_path}': {e}")
+            raise
+
+        # Call the PIL-based prediction method
+        return self.predict_from_pil(image)
+
+
+def info():
+    _script_info(__all__)