/ludwig/features/image_feature.py

#! /usr/bin/env python
# Copyright (c) 2023 Predibase, Inc., 2019 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import logging
import os
import warnings
from collections import Counter
from collections.abc import Callable
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torchvision import transforms
from torchvision.transforms import functional as F
from torchvision.transforms.functional import normalize

from ludwig.constants import (
    COLUMN,
    ENCODER,
    HEIGHT,
    IMAGE,
    IMAGENET1K,
    INFER_IMAGE_DIMENSIONS,
    INFER_IMAGE_MAX_HEIGHT,
    INFER_IMAGE_MAX_WIDTH,
    INFER_IMAGE_NUM_CLASSES,
    INFER_IMAGE_SAMPLE_SIZE,
    LOGITS,
    NAME,
    NUM_CHANNELS,
    PREDICTIONS,
    PREPROCESSING,
    PROC_COLUMN,
    REQUIRES_EQUAL_DIMENSIONS,
    SRC,
    TYPE,
    WIDTH,
)
from ludwig.data.lazy_utils import resolve_lazy_cache_dir
from ludwig.encoders.base import Encoder
from ludwig.encoders.image.torchvision import TVModelVariant
from ludwig.features.base_feature import (
    BasePostprocessingModule,
    BasePreprocessingModule,
    FeaturePreprocessingMixin,
    InputFeature,
    OutputFeature,
    PredictModule,
)
from ludwig.schema.features.augmentation.base import BaseAugmentationConfig
from ludwig.schema.features.augmentation.image import (
    AutoAugmentationConfig,
    RandomBlurConfig,
    RandomBrightnessConfig,
    RandomContrastConfig,
    RandomHorizontalFlipConfig,
    RandomRotateConfig,
    RandomVerticalFlipConfig,
)
from ludwig.schema.features.image_feature import ImageInputFeatureConfig, ImageOutputFeatureConfig
from ludwig.types import (
    FeatureMetadataDict,
    FeaturePostProcessingOutputDict,
    ModelConfigDict,
    PreprocessingConfigDict,
    TrainingSetMetadataDict,
)
from ludwig.utils import output_feature_utils
from ludwig.utils.augmentation_utils import get_augmentation_op, register_augmentation_op
from ludwig.utils.data_utils import get_abs_path
from ludwig.utils.dataframe_utils import is_dask_series_or_df
from ludwig.utils.fs_utils import has_remote_protocol
from ludwig.utils.image_utils import (
    get_class_mask_from_image,
    get_gray_default_image,
    get_image_from_class_mask,
    get_unique_channels,
    grayscale,
    num_channels_in_image,
    read_image_from_bytes_obj,
    read_image_from_path,
    resize_image,
    ResizeChannels,
    torchvision_model_registry,
)
from ludwig.utils.misc_utils import set_default_value
from ludwig.utils.types import PreprocessingInput, Series

# constants used for Ludwig image preprocessing
IMAGENET1K_MEAN = [0.485, 0.456, 0.406]
IMAGENET1K_STD = [0.229, 0.224, 0.225]


logger = logging.getLogger(__name__)


def _extract_image_path(entry: object) -> str | None:
    """Return a usable on-disk path from a PIL Image or dict entry, or ``None``.

    PIL Images opened from disk retain their source path in the ``.filename``
    attribute.  HuggingFace ``Image`` columns may deliver dicts that contain a
    ``"path"`` key pointing to the cached file.

    Parameters
    ----------
    entry:
        A single image entry from a dataset column.  May be a ``PIL.Image.Image``
        with a ``.filename`` attribute, a dict with a ``"path"`` key, or any
        other type.

    Returns
    -------
    str | None
        The existing on-disk path when one is found and the file is present on
        disk; ``None`` otherwise.
    """
    # PIL Image opened from a file retains its filename
    filename = getattr(entry, "filename", None)
    if filename and isinstance(filename, str) and os.path.isfile(filename):
        return filename

    # HuggingFace-style dict with a "path" key
    if isinstance(entry, dict):
        path = entry.get("path")
        if path and isinstance(path, str) and os.path.isfile(path):
            return path

    return None


def _cache_image_column_to_disk(
    column,
    cache_dir: Path,
    feature_name: str,
) -> list[str]:
    """Write every image entry in *column* to PNG files in *cache_dir*.

    Entries that already have an accessible on-disk path (PIL ``.filename`` or
    dict ``"path"``) are reused without copying.  All other entries (in-memory
    PIL Images, raw bytes, or numpy arrays) are saved as PNG files.  Files
    that already exist in *cache_dir* are not overwritten (idempotent
    behaviour), which makes re-runs after a crash safe.

    Parameters
    ----------
    column:
        An iterable of image entries.  Each entry may be:

        * a ``PIL.Image.Image`` — with or without a ``.filename`` attribute;
        * ``bytes`` containing a valid encoded image;
        * a ``numpy.ndarray`` with shape ``(H, W)``, ``(H, W, C)``, or
          ``(C, H, W)`` and dtype ``uint8``;
        * a dict with at least a ``"bytes"`` or ``"path"`` key (HuggingFace
          ``Image`` column format).
    cache_dir:
        Directory in which to write PNG files.  Must already exist.
    feature_name:
        Used as a filename prefix so that files are recognisable during
        debugging.

    Returns
    -------
    list[str]
        Ordered list of absolute paths (one per entry), suitable for direct
        use as a ``LazyColumn`` path list.

    Raises
    ------
    ValueError
        If an entry has an unrecognised type that cannot be converted to a PIL
        Image for saving.
    """
    # Lazy import — PIL is an optional dependency; fail loudly only when used.
    from PIL import Image as PILImage

    paths: list[str] = []
    for idx, entry in enumerate(column):
        # --- Try to reuse an existing on-disk path ---
        existing = _extract_image_path(entry)
        if existing is not None:
            paths.append(existing)
            continue

        # --- Convert entry to a PIL Image ---
        if isinstance(entry, PILImage.Image):
            pil_img = entry
        elif isinstance(entry, bytes):
            import io

            pil_img = PILImage.open(io.BytesIO(entry)).copy()
        elif isinstance(entry, np.ndarray):
            # Handle (C, H, W) → (H, W, C) for RGB
            if entry.ndim == 3 and entry.shape[0] in (1, 3, 4):
                entry = np.transpose(entry, (1, 2, 0))
            pil_img = PILImage.fromarray(entry.astype(np.uint8))
        elif isinstance(entry, dict):
            # HuggingFace Image column format: {"bytes": <bytes>, "path": <str|None>}
            raw_bytes = entry.get("bytes")
            if raw_bytes:
                import io

                pil_img = PILImage.open(io.BytesIO(raw_bytes)).copy()
            else:
                raise ValueError(
                    f"Image entry [{idx}] in feature '{feature_name}' is a dict but contains "
                    "neither a usable 'path' nor 'bytes' key."
                )
        else:
            # Try torch.Tensor as a last resort (optional dependency)
            try:
                import torch

                if isinstance(entry, torch.Tensor):
                    arr = entry.detach().cpu().numpy()
                    # Handle (C, H, W) → (H, W, C) for colour tensors
                    if arr.ndim == 3 and arr.shape[0] in (1, 3, 4):
                        arr = np.transpose(arr, (1, 2, 0))
                    pil_img = PILImage.fromarray(arr.astype(np.uint8))
                else:
                    raise TypeError
            except (ImportError, TypeError):
                raise ValueError(
                    f"Image entry [{idx}] in feature '{feature_name}' has unrecognised "
                    f"type {type(entry).__name__!r}.  "
                    "Expected PIL.Image, bytes, numpy.ndarray, torch.Tensor, or dict."
                )

        dest_path = str(cache_dir / f"{feature_name}_{idx:08d}.png")
        if not os.path.isfile(dest_path):
            pil_img.save(dest_path, format="PNG")
        paths.append(dest_path)

    return paths


###
# Image specific augmentation operations
###
@register_augmentation_op(name="auto_augmentation", features=IMAGE)
class AutoAugment(torch.nn.Module):
    def __init__(self, config: AutoAugmentationConfig):
        super().__init__()
        self.auto_augmentation_method = config.method
        self.augmentation_method = self.get_augmentation_method()

    def get_augmentation_method(self):
        if self.auto_augmentation_method == "trivial_augment":
            return transforms.TrivialAugmentWide()
        if self.auto_augmentation_method == "auto_augment":
            return transforms.AutoAugment()
        if self.auto_augmentation_method == "rand_augment":
            return transforms.RandAugment()
        raise ValueError(f"Unsupported auto-augmentation method: {self.auto_augmentation_method}")

    def forward(self, imgs: torch.Tensor) -> torch.Tensor:
        method = self.augmentation_method
        uint8imgs = imgs.to(torch.uint8)
        augmented_imgs = method(uint8imgs)

        return augmented_imgs.to(torch.float32)


@register_augmentation_op(name="random_vertical_flip", features=IMAGE)
class RandomVFlip(torch.nn.Module):
    def __init__(
        self,
        config: RandomVerticalFlipConfig,
    ):
        super().__init__()

    def forward(self, imgs):
        if torch.rand(1) < 0.5:
            imgs = F.vflip(imgs)

        return imgs


@register_augmentation_op(name="random_horizontal_flip", features=IMAGE)
class RandomHFlip(torch.nn.Module):
    def __init__(
        self,
        config: RandomHorizontalFlipConfig,
    ):
        super().__init__()

    def forward(self, imgs):
        if torch.rand(1) < 0.5:
            imgs = F.hflip(imgs)

        return imgs


@register_augmentation_op(name="random_rotate", features=IMAGE)
class RandomRotate(torch.nn.Module):
    def __init__(self, config: RandomRotateConfig):
        super().__init__()
        self.degree = config.degree

    def forward(self, imgs):
        if torch.rand(1) < 0.5:
            # map angle to interval (-degree, +degree)
            angle = (torch.rand(1) * 2 * self.degree - self.degree).item()
            return F.rotate(imgs, angle)
        else:
            return imgs


@register_augmentation_op(name="random_contrast", features=IMAGE)
class RandomContrast(torch.nn.Module):
    def __init__(self, config: RandomContrastConfig):
        super().__init__()
        self.min_contrast = config.min
        self.contrast_adjustment_range = config.max - config.min

    def forward(self, imgs):
        if torch.rand(1) < 0.5:
            # random contrast adjustment
            adjust_factor = (torch.rand(1) * self.contrast_adjustment_range + self.min_contrast).item()
            return F.adjust_contrast(imgs, adjust_factor)
        else:
            return imgs


@register_augmentation_op(name="random_brightness", features=IMAGE)
class RandomBrightness(torch.nn.Module):
    def __init__(self, config: RandomBrightnessConfig):
        super().__init__()
        self.min_brightness = config.min
        self.brightness_adjustment_range = config.max - config.min

    def forward(self, imgs):
        if torch.rand(1) < 0.5:
            # random contrast adjustment
            adjust_factor = (torch.rand(1) * self.brightness_adjustment_range + self.min_brightness).item()
            return F.adjust_brightness(imgs, adjust_factor)
        else:
            return imgs


@register_augmentation_op(name="random_blur", features=IMAGE)
class RandomBlur(torch.nn.Module):
    def __init__(self, config: RandomBlurConfig):
        super().__init__()
        self.kernel_size = [config.kernel_size, config.kernel_size]

    def forward(self, imgs):
        if torch.rand(1) < 0.5:
            imgs = F.gaussian_blur(imgs, self.kernel_size)

        return imgs


class ImageAugmentation(torch.nn.Module):
    def __init__(
        self,
        augmentation_list: list[BaseAugmentationConfig],
        normalize_mean: list[float] | None = None,
        normalize_std: list[float] | None = None,
    ):
        super().__init__()

        logger.debug(f"Creating augmentation pipeline: {augmentation_list}")

        self.normalize_mean = normalize_mean
        self.normalize_std = normalize_std

        if self.training:
            self.augmentation_steps = torch.nn.Sequential()
            for aug_config in augmentation_list:
                try:
                    aug_op = get_augmentation_op(IMAGE, aug_config.type)
                    self.augmentation_steps.append(aug_op(aug_config))
                except KeyError:
                    raise ValueError(f"Invalid augmentation operation specification: {aug_config}")
        else:
            self.augmentation_steps = None

    def forward(self, imgs):
        if self.augmentation_steps:
            # convert from float to uint8 values - this is required for the augmentation
            imgs = self._convert_back_to_uint8(imgs)

            logger.debug("Executing augmentation pipeline steps: %s", self.augmentation_steps)
            imgs = self.augmentation_steps(imgs)

            # convert back to float32 values and renormalize if needed
            imgs = self._renormalize_image(imgs)

        return imgs

    # function to partially undo the TorchVision ImageClassification transformation.
    #  back out the normalization step and convert from float32 to uint8 dtype
    #  to make the tensor displayable as an image
    #  crop size remains the same
    def _convert_back_to_uint8(self, images):
        if self.normalize_mean:
            mean = torch.as_tensor(self.normalize_mean, dtype=torch.float32).view(-1, 1, 1)
            std = torch.as_tensor(self.normalize_std, dtype=torch.float32).view(-1, 1, 1)
            return images.mul(std).add(mean).mul(255.0).type(torch.uint8)
        else:
            return images.mul(255.0).type(torch.uint8)

    # function to redo part of the TorchVision ImageClassification transformation.
    #  convert uint8 to float32
    #  apply the imagenet1k normalization
    def _renormalize_image(self, images):
        if self.normalize_mean:
            mean = torch.as_tensor(self.normalize_mean, dtype=torch.float32).view(-1, 1, 1)
            std = torch.as_tensor(self.normalize_std, dtype=torch.float32).view(-1, 1, 1)
            return images.type(torch.float32).div(255.0).sub(mean).div(std)
        else:
            return images.type(torch.float32).div(255.0)


@dataclass
class ImageTransformMetadata:
    height: int
    width: int
    num_channels: int


def _get_torchvision_transform(
    torchvision_parameters: TVModelVariant,
) -> tuple[torch.nn.Module, ImageTransformMetadata]:
    """Returns a torchvision transform that is compatible with the model variant.

    Note that the raw torchvision transform is not returned. Instead, a Sequential module that includes
    image resizing is returned. This is because the raw torchvision transform assumes that the input image has
    three channels, which is not always the case with images input into Ludwig.

    Args:
        torchvision_parameters: The parameters for the torchvision model variant.
    Returns:
        (torchvision_transform, transform_metadata): A torchvision transform and the metadata for the transform.
    """
    torchvision_transform_raw = torchvision_parameters.model_weights.DEFAULT.transforms()
    torchvision_transform = torch.nn.Sequential(
        ResizeChannels(num_channels=3),
        torchvision_transform_raw,
    )
    transform_metadata = ImageTransformMetadata(
        height=torchvision_transform_raw.crop_size[0],
        width=torchvision_transform_raw.crop_size[0],
        num_channels=len(torchvision_transform_raw.mean),
    )
    return (torchvision_transform, transform_metadata)


def _get_torchvision_parameters(model_type: str, model_variant: str) -> TVModelVariant:
    return torchvision_model_registry.get(model_type).get(model_variant)


def is_torchvision_encoder(encoder_obj: Encoder) -> bool:
    # TODO(travis): do this through an interface rather than conditional logic
    from ludwig.encoders.image.torchvision import TVBaseEncoder

    return isinstance(encoder_obj, TVBaseEncoder)


class _ImagePreprocessing(BasePreprocessingModule):
    """Torchscript-enabled version of preprocessing done by ImageFeatureMixin.add_feature_data."""

    def __init__(
        self,
        metadata: TrainingSetMetadataDict,
        torchvision_transform: torch.nn.Module | None = None,
        transform_metadata: ImageTransformMetadata | None = None,
    ):
        super().__init__()

        self.resize_method = metadata["preprocessing"]["resize_method"]
        self.torchvision_transform = torchvision_transform
        if transform_metadata is not None:
            self.height = transform_metadata.height
            self.width = transform_metadata.width
            self.num_channels = transform_metadata.num_channels
            self.channel_class_map = torch.Tensor([])
        else:
            self.height = metadata["preprocessing"]["height"]
            self.width = metadata["preprocessing"]["width"]
            self.num_channels = metadata["preprocessing"]["num_channels"]
            self.channel_class_map = torch.ByteTensor(metadata["preprocessing"]["channel_class_map"])

    def forward(self, v: PreprocessingInput) -> torch.Tensor:
        """Takes a list of images and adjusts the size and number of channels as specified in the metadata.

        If `v` is already a torch.Tensor, we assume that the images are already preprocessed to be the same size.
        """
        # Nested conditional is a workaround to short-circuit boolean evaluation.
        if not torch.jit.isinstance(v, list[torch.Tensor]):
            if not torch.jit.isinstance(v, torch.Tensor):
                raise ValueError(f"Unsupported input: {v}")

        if self.torchvision_transform is not None:
            # perform pre-processing for torchvision pretrained model encoders
            if torch.jit.isinstance(v, list[torch.Tensor]):
                imgs = [self.torchvision_transform(img) for img in v]
            else:
                # convert batch of image tensors to a list and then run torchvision pretrained
                # model transforms on each image
                imgs = [self.torchvision_transform(img) for img in torch.unbind(v)]

            # collect the list of images into a batch
            imgs_stacked = torch.stack(imgs)
        else:
            # perform pre-processing for Ludwig defined image encoders
            if torch.jit.isinstance(v, list[torch.Tensor]):
                imgs = [resize_image(img, (self.height, self.width), self.resize_method) for img in v]
                imgs_stacked = torch.stack(imgs)
            else:
                imgs_stacked = v

            _, num_channels, height, width = imgs_stacked.shape

            # Ensure images are the size expected by the model
            if height != self.height or width != self.width:
                imgs_stacked = resize_image(imgs_stacked, (self.height, self.width), self.resize_method)

            # Ensures images have the number of channels expected by the model
            if num_channels != self.num_channels:
                if self.num_channels == 1:
                    imgs_stacked = grayscale(imgs_stacked)
                elif num_channels < self.num_channels:
                    extra_channels = self.num_channels - num_channels
                    imgs_stacked = torch.nn.functional.pad(imgs_stacked, [0, 0, 0, 0, 0, extra_channels])
                else:
                    raise ValueError(
                        f"Number of channels cannot be reconciled. metadata.num_channels = "
                        f"{self.num_channels}, but imgs.shape[1] = {num_channels}"
                    )

            # Create class-masked images if required
            if self.channel_class_map.shape[0]:
                masks = []
                for img in imgs_stacked:
                    mask = get_class_mask_from_image(self.channel_class_map, img)
                    masks.append(mask)
                imgs_stacked = torch.stack(masks)
            else:
                imgs_stacked = imgs_stacked.type(torch.float32) / 255

        return imgs_stacked


class _ImagePostprocessing(BasePostprocessingModule):
    def __init__(self):
        super().__init__()
        self.logits_key = LOGITS
        self.predictions_key = PREDICTIONS

    def forward(self, preds: dict[str, torch.Tensor], feature_name: str) -> FeaturePostProcessingOutputDict:
        predictions = output_feature_utils.get_output_feature_tensor(preds, feature_name, self.predictions_key)
        logits = output_feature_utils.get_output_feature_tensor(preds, feature_name, self.logits_key)

        return {self.predictions_key: predictions, self.logits_key: logits}


class _ImagePredict(PredictModule):
    def forward(self, inputs: dict[str, torch.Tensor], feature_name: str) -> dict[str, torch.Tensor]:
        predictions = output_feature_utils.get_output_feature_tensor(inputs, feature_name, self.predictions_key)
        logits = output_feature_utils.get_output_feature_tensor(inputs, feature_name, self.logits_key)

        return {self.predictions_key: predictions, self.logits_key: logits}


class ImageFeatureMixin(FeaturePreprocessingMixin):
    @staticmethod
    def type():
        return IMAGE

    @staticmethod
    def cast_column(column, backend):
        return column

    @staticmethod
    def get_feature_meta(
        config: ModelConfigDict,
        column,
        preprocessing_parameters: PreprocessingConfigDict,
        backend,
        is_input_feature: bool,
    ) -> FeatureMetadataDict:
        return {PREPROCESSING: preprocessing_parameters}

    @staticmethod
    def _read_image_if_bytes_obj_and_resize(
        img_entry: bytes | torch.Tensor | np.ndarray | str,
        img_width: int,
        img_height: int,
        should_resize: bool,
        num_channels: int,
        resize_method: str,
        user_specified_num_channels: bool,
        standardize_image: str,
        channel_class_map: torch.Tensor,
    ) -> np.ndarray | None:
        """Helper method to read and resize an image according to model definition. If the user doesn't specify a
        number of channels, we use the first image in the dataset as the source of truth. If any image in the
        dataset doesn't have the same number of channels as the first image, raise an exception.

        If the user specifies a number of channels, we try to convert all the images to the specifications by
        dropping channels/padding 0 channels.

        Args:
            img_entry: if str, file path to the image; otherwise a torch.Tensor, np.ndarray, or bytes of the image.
            img_width: expected width of the image.
            img_height: expected height of the image.
            should_resize: Should the image be resized?
            resize_method: type of resizing method.
            num_channels: expected number of channels in the first image.
            user_specified_num_channels: did the user specify num channels?
            standardize_image: specifies whether to standardize image with imagenet1k specifications.
            channel_class_map: A tensor mapping channel values to classes, where dim=0 is the class.

        Returns:
            image object as a numpy array.
        """

        if isinstance(img_entry, bytes):
            img = read_image_from_bytes_obj(img_entry, num_channels)
        elif isinstance(img_entry, str):
            img = read_image_from_path(img_entry, num_channels)
        elif isinstance(img_entry, np.ndarray):
            img = torch.from_numpy(np.array(img_entry, copy=True)).permute(2, 0, 1)
        else:
            img = img_entry

        if not isinstance(img, torch.Tensor):
            warnings.warn(f"Image with value {img} cannot be read")
            return None

        img_num_channels = num_channels_in_image(img)
        # Convert to grayscale if needed.
        if num_channels == 1 and img_num_channels != 1:
            img = grayscale(img)
            img_num_channels = 1

        if should_resize:
            img = resize_image(img, (img_height, img_width), resize_method)

        if user_specified_num_channels:
            # Number of channels is specified by the user
            # img_padded = np.zeros((img_height, img_width, num_channels),
            #                       dtype=np.uint8)
            # min_num_channels = min(num_channels, img_num_channels)
            # img_padded[:, :, :min_num_channels] = img[:, :, :min_num_channels]
            # img = img_padded
            if num_channels > img_num_channels:
                extra_channels = num_channels - img_num_channels
                img = torch.nn.functional.pad(img, [0, 0, 0, 0, 0, extra_channels])

            if img_num_channels != num_channels:
                logger.warning(
                    f"Image has {img_num_channels} channels, where as {num_channels} "
                    "channels are expected. Dropping/adding channels "
                    "with 0s as appropriate"
                )
        else:
            # If the image isn't like the first image, raise exception
            if img_num_channels != num_channels:
                raise ValueError(
                    f"Image has {img_num_channels} channels, unlike the first image, which "
                    f"has {num_channels} channels. Make sure all the images have the same "
                    "number of channels or use the num_channels property in "
                    "image preprocessing"
                )

        if img.shape[1] != img_height or img.shape[2] != img_width:
            raise ValueError(
                "Images are not of the same size. "
                f"Expected size is {[img_height, img_width, num_channels]}, "
                f"current image size is {img.shape}."
                "Images are expected to be all of the same size "
                "or explicit image width and height are expected "
                "to be provided. "
                "Additional information: "
                "https://ludwig-ai.github.io/ludwig-docs/latest/configuration/features/image_features"
                "#image-features-preprocessing"
            )

        # Create class-masked image if required
        if channel_class_map.shape[0]:
            img = get_class_mask_from_image(channel_class_map, img)
        else:
            # casting and rescaling
            img = img.type(torch.float32) / 255

            if standardize_image == IMAGENET1K:
                img = normalize(img, mean=IMAGENET1K_MEAN, std=IMAGENET1K_STD)

        return img.numpy()

    @staticmethod
    def _read_image_with_pretrained_transform(
        img_entry: bytes | torch.Tensor | np.ndarray,
        transform_fn: Callable,
    ) -> np.ndarray | None:
        if isinstance(img_entry, bytes):
            img = read_image_from_bytes_obj(img_entry)
        elif isinstance(img_entry, str):
            img = read_image_from_path(img_entry)
        elif isinstance(img_entry, np.ndarray):
            img = torch.from_numpy(img_entry).permute(2, 0, 1)
        else:
            img = img_entry

        if not isinstance(img, torch.Tensor):
            warnings.warn(f"Image with value {img} cannot be read")
            return None

        img = transform_fn(img)

        return img.numpy()

    @staticmethod
    def _make_lazy_decode_fn(
        img_width: int,
        img_height: int,
        should_resize: bool,
        num_channels: int,
        resize_method: str,
        user_specified_num_channels: bool,
        standardize_image: str,
        channel_class_map,
        default_image: np.ndarray,
    ):
        """Return a per-path decode function suitable for use in ``LazyColumn``.

        The returned callable is stateless (captures only plain values / arrays)
        so it is safe to share across DataLoader worker threads.
        """
        _resize_fn = partial(
            ImageFeatureMixin._read_image_if_bytes_obj_and_resize,
            img_width=img_width,
            img_height=img_height,
            should_resize=should_resize,
            num_channels=num_channels,
            resize_method=resize_method,
            user_specified_num_channels=user_specified_num_channels,
            standardize_image=standardize_image,
            channel_class_map=channel_class_map,
        )

        def decode(path_or_bytes) -> np.ndarray:
            result = _resize_fn(path_or_bytes)
            if not isinstance(result, np.ndarray):
                return default_image
            return result

        decode.__name__ = "image_lazy_decode"
        return decode

    @staticmethod
    def _set_image_and_height_equal_for_encoder(
        width: int, height: int, preprocessing_parameters: dict, encoder_type: str
    ) -> tuple[int, int]:
        """Some pretrained image encoders require images with the same dimension, or images with a specific width
        and heigh values. The returned width and height are set based on compatibility with the downstream encoder
        using the encoder parameters for the feature.

        Args:
            width: Represents the width of the image. This is either specified in the user config, or inferred using
                a sample of images.
            height: Represents the height of the image. This is either specified in the user config, or inferred using
                a sample of images.
            preprocessing_parameters: Parameters defining how the image feature should be preprocessed
            encoder_type: The name of the encoder

        Return:
            (width, height) Updated width and height so that they are equal
        """

        if preprocessing_parameters[REQUIRES_EQUAL_DIMENSIONS] and height != width:
            width = height = min(width, height)
            # Update preprocessing parameters dictionary to reflect new height and width values
            preprocessing_parameters["width"] = width
            preprocessing_parameters["height"] = height
            logger.info(f"Set image feature height and width to {width} to be compatible with {encoder_type} encoder.")
        return width, height

    @staticmethod
    def _infer_image_size(
        image_sample: list[torch.Tensor],
        max_height: int,
        max_width: int,
        preprocessing_parameters: dict,
        encoder_type: str,
    ) -> tuple[int, int]:
        """Infers the size to use from a group of images. The returned height will be the average height of images
        in image_sample rounded to the nearest integer, or max_height. Likewise for width.

        Args:
            image_sample: Sample of images to use to infer image size. Must be formatted as [channels, height, width].
            max_height: Maximum height.
            max_width: Maximum width.
            preprocessing_parameters: Parameters defining how the image feature should be preprocessed
            encoder_type: The name of the encoder

        Return:
            (height, width) The inferred height and width.
        """

        height_avg = sum(x.shape[1] for x in image_sample) / len(image_sample)
        width_avg = sum(x.shape[2] for x in image_sample) / len(image_sample)
        height = min(int(round(height_avg)), max_height)
        width = min(int(round(width_avg)), max_width)

        # Update height and width if the downstream encoder requires images
        # with  the same dimension or specific width and height values
        width, height = ImageFeatureMixin._set_image_and_height_equal_for_encoder(
            width, height, preprocessing_parameters, encoder_type
        )

        logger.debug(f"Inferring height: {height} and width: {width}")
        return height, width

    @staticmethod
    def _infer_number_of_channels(image_sample: list[torch.Tensor]):
        """Infers the channel depth to use from a group of images.

        We make the assumption that the majority of datasets scraped from the web will be RGB, so if we get a mixed bag
        of images we should default to that. However, if the majority of the sample images have a specific channel depth
        (other than 3) this is probably intentional so we keep it, but log an info message.
        """
        n_images = len(image_sample)
        channel_frequency = Counter([num_channels_in_image(x) for x in image_sample])
        if channel_frequency[1] > n_images / 2:
            # If the majority of images in sample are 1 channel, use 1.
            num_channels = 1
        elif channel_frequency[2] > n_images / 2:
            # If the majority of images in sample are 2 channel, use 2.
            num_channels = 2
        elif channel_frequency[4] > n_images / 2:
            # If the majority of images in sample are 4 channel, use 4.
            num_channels = 4
        else:
            # Default case: use 3 channels.
            num_channels = 3
        logger.info(f"Inferring num_channels from the first {n_images} images.")
        logger.info("\n".join([f"  images with {k} channels: {v}" for k, v in sorted(channel_frequency.items())]))
        if num_channels == max(channel_frequency, key=channel_frequency.get):
            logger.info(
                f"Using {num_channels} channels because it is the majority in sample. If an image with"
                f" a different depth is read, will attempt to convert to {num_channels} channels."
            )
        else:
            logger.info(f"Defaulting to {num_channels} channels.")
        logger.info(
            "To explicitly set the number of channels, define num_channels in the preprocessing dictionary of "
            "the image input feature config."
        )
        return num_channels

    @staticmethod
    def _infer_image_num_classes(
        image_sample: list[torch.Tensor],
        num_channels: int,
        num_classes: int,
    ) -> torch.Tensor:
        """Infers the number of channel classes from a group of images (for image segmentation). The returned
        tensor contains the channel value for each class, where dim=0 is the class.

        Args:
            image_sample: Sample of images to use to infer image size. Must be formatted as [channels, height, width].
            num_channels: Expected number of channels
            num_classes: Expected number of channel classes or None

        Return:
            channel_class_map: A tensor mapping channel values to classes, where dim=0 is the class.
        """
        n_images = len(image_sample)
        logger.info(f"Inferring num_classes from the first {n_images} images.")
        channel_class_map = get_unique_channels(image_sample, num_channels, num_classes)

        inferred_num_classes = channel_class_map.shape[0]
        if num_classes:
            if num_classes < inferred_num_classes:
                raise ValueError(
                    f"Images inferred num classes {inferred_num_classes} exceeds `num_classes` {num_classes}."
                )
            elif num_classes > inferred_num_classes:
                logger.warning(
                    f"Images inferred num classes {inferred_num_classes} does not match `num_classes` {num_classes}. "
                    f"Using inferred num classes {inferred_num_classes}."
                )

        return channel_class_map

    @staticmethod
    def _finalize_preprocessing_parameters(
        preprocessing_parameters: dict,
        encoder_type: str,
        column: Series,
    ) -> tuple:
        """Helper method to determine the height, width and number of channels for preprocessing the image data.

        This is achieved by looking at the parameters provided by the user. When there are some missing parameters, we
        fall back on to the first image in the dataset. The assumption being that all the images in the data are
        expected be of the same size with the same number of channels.

        Args:
            preprocessing_parameters: Parameters defining how the image feature should be preprocessed
            encoder_type: The name of the encoder
            column: The data itself. Can be a Pandas, Modin or Dask series.
        """

        explicit_height_width = preprocessing_parameters[HEIGHT] or preprocessing_parameters[WIDTH]
        explicit_num_channels = NUM_CHANNELS in preprocessing_parameters and preprocessing_parameters[NUM_CHANNELS]

        if preprocessing_parameters[INFER_IMAGE_DIMENSIONS] and not (explicit_height_width and explicit_num_channels):
            sample_size = min(len(column), preprocessing_parameters[INFER_IMAGE_SAMPLE_SIZE])
        else:
            sample_size = 1  # Take first image

        sample = []
        sample_num_bytes = []
        failed_entries = []
        for image_entry in column.head(sample_size):
            if isinstance(image_entry, bytes):
                image = read_image_from_bytes_obj(image_entry)
            elif isinstance(image_entry, str):
                # Tries to read image as PNG or numpy file from the path.
                image, num_bytes = read_image_from_path(image_entry, return_num_bytes=True)
                if num_bytes is not None:
                    sample_num_bytes.append(num_bytes)
            else:
                image = image_entry

            if isinstance(image, torch.Tensor):
                sample.append(image)
            elif isinstance(image, np.ndarray):
                sample.append(torch.from_numpy(image).permute(2, 0, 1))
            else:
                failed_entries.append(image_entry)
        if len(sample) == 0:
            failed_entries_repr = "\n\t- ".join(failed_entries)
            raise ValueError(
                f"Images dimensions cannot be inferred. Failed to read {sample_size} images as samples:"
                f"\n\t- {failed_entries_repr}."
            )

        should_resize = False
        if explicit_height_width:
            should_resize = True
            try:
                height = int(preprocessing_parameters[HEIGHT])
                width = int(preprocessing_parameters[WIDTH])
                # Update height and width if the downstream encoder requires images
                # with the same dimension or specific width and height values
                width, height = ImageFeatureMixin._set_image_and_height_equal_for_encoder(
                    width, height, preprocessing_parameters, encoder_type
                )
            except ValueError as e:
                raise ValueError("Image height and width must be set and have positive integer values: " + str(e))
            if height <= 0 or width <= 0:
                raise ValueError("Image height and width must be positive integers")
        else:
            # User hasn't specified height and width.
            # Default to inferring from sample or first image.
            if preprocessing_parameters[INFER_IMAGE_DIMENSIONS]:
                should_resize = True
                height, width = ImageFeatureMixin._infer_image_size(
                    sample,
                    max_height=preprocessing_parameters[INFER_IMAGE_MAX_HEIGHT],
                    max_width=preprocessing_parameters[INFER_IMAGE_MAX_WIDTH],
                    preprocessing_parameters=preprocessing_parameters,
                    encoder_type=encoder_type,
                )
            else:
                raise ValueError(
                    "Explicit image width/height are not set, infer_image_dimensions is false, "
                    "and first image cannot be read, so image dimensions are unknown"
                )

        if explicit_num_channels:
            # User specified num_channels in the model/feature config
            user_specified_num_channels = True
            num_channels = preprocessing_parameters[NUM_CHANNELS]
        else:
            user_specified_num_channels = False
            if preprocessing_parameters[INFER_IMAGE_DIMENSIONS]:
                user_specified_num_channels = True
                num_channels = ImageFeatureMixin._infer_number_of_channels(sample)
            elif len(sample) > 0:
                num_channels = num_channels_in_image(sample[0])
            else:
                raise ValueError(
                    "Explicit image num channels is not set, infer_image_dimensions is false, "
                    "and first image cannot be read, so image num channels is unknown"
                )

        if not isinstance(num_channels, int):
            raise ValueError(f"Number of image channels needs to be an integer, got {type(num_channels).__name__}.")

        average_file_size = np.mean(sample_num_bytes) if sample_num_bytes else None

        standardize_image = preprocessing_parameters["standardize_image"]
        if standardize_image == "imagenet1k" and num_channels != 3:
            warnings.warn(
                f"'standardize_image=imagenet1k' is defined only for 'num_channels=3' but "
                f"detected 'num_channels={num_channels}'.  For this situation setting 'standardize_image=None'.",
                RuntimeWarning,
            )
            standardize_image = None

        if preprocessing_parameters[INFER_IMAGE_NUM_CLASSES] or preprocessing_parameters["num_classes"]:
            channel_class_map = ImageFeatureMixin._infer_image_num_classes(
                sample, num_channels, preprocessing_parameters["num_classes"]
            )
        else:
            channel_class_map = torch.Tensor([])

        return (
            should_resize,
            width,
            height,
            num_channels,
            user_specified_num_channels,
            average_file_size,
            standardize_image,
            channel_class_map,
        )

    @staticmethod
    def add_feature_data(
        feature_config,
        input_df,
        proc_df,
        metadata,
        preprocessing_parameters: PreprocessingConfigDict,
        backend,
        skip_save_processed_input,
    ):
        set_default_value(feature_config[PREPROCESSING], "in_memory", preprocessing_parameters["in_memory"])

        name = feature_config[NAME]
        column = input_df[feature_config[COLUMN]]
        encoder_type = feature_config[ENCODER][TYPE] if ENCODER in feature_config else None

        src_path = None
        if SRC in metadata:
            src_path = os.path.dirname(os.path.abspath(metadata.get(SRC)))
        abs_path_column = backend.df_engine.map_objects(
            column,
            lambda row: get_abs_path(src_path, row) if isinstance(row, str) and not has_remote_protocol(row) else row,
        )

        # determine if specified encoder is a torchvision model
        model_type = feature_config[ENCODER].get("type", None) if ENCODER in feature_config else None
        model_variant = feature_config[ENCODER].get("model_variant") if ENCODER in feature_config else None
        if model_variant:
            torchvision_parameters = _get_torchvision_parameters(model_type, model_variant)
        else:
            torchvision_parameters = None

        if torchvision_parameters:
            logger.warning(
                f"Using the transforms specified for the torchvision model {model_type} {model_variant} "
                f"This includes setting the number of channels is 3 and resizing the image to the needs of the model."
            )

            torchvision_transform, transform_metadata = _get_torchvision_transform(torchvision_parameters)

            # torchvision_parameters is not None
            # perform torchvision model transformations
            read_image_if_bytes_obj_and_resize = partial(
                ImageFeatureMixin._read_image_with_pretrained_transform,
                transform_fn=torchvision_transform,
            )
            average_file_size = None

            # save weight specification in preprocessing section
            preprocessing_parameters["torchvision_model_default_weights"] = (
                f"{torchvision_parameters.model_weights.DEFAULT}"
            )

            # add torchvision model id to preprocessing section for torchscript
            preprocessing_parameters["torchvision_model_type"] = model_type
            preprocessing_parameters["torchvision_model_variant"] = model_variant

            # get required setup parameters for in_memory = False processing
            height = transform_metadata.height
            width = transform_metadata.width
            num_channels = transform_metadata.num_channels
            channel_class_map = torch.Tensor([])
        else:
            # torchvision_parameters is None
            # perform Ludwig specified transformations
            (
                should_resize,
                width,
                height,
                num_channels,
                user_specified_num_channels,
                average_file_size,
                standardize_image,
                channel_class_map,
            ) = ImageFeatureMixin._finalize_preprocessing_parameters(
                preprocessing_parameters, encoder_type, abs_path_column
            )

            metadata[name][PREPROCESSING]["height"] = height
            metadata[name][PREPROCESSING]["width"] = width
            metadata[name][PREPROCESSING]["num_channels"] = num_channels
            metadata[name][PREPROCESSING]["num_classes"] = channel_class_map.shape[0]
            metadata[name][PREPROCESSING]["channel_class_map"] = channel_class_map.tolist()

            read_image_if_bytes_obj_and_resize = partial(
                ImageFeatureMixin._read_image_if_bytes_obj_and_resize,
                img_width=width,
                img_height=height,
                should_resize=should_resize,
                num_channels=num_channels,
                resize_method=preprocessing_parameters["resize_method"],
                user_specified_num_channels=user_specified_num_channels,
                standardize_image=standardize_image,
                channel_class_map=channel_class_map,
            )

        # TODO: alternatively use get_average_image() for unreachable images
        if channel_class_map.shape[0]:
            default_image = get_gray_default_image(1, height, width).squeeze(0)
            metadata[name]["reshape"] = (height, width)
        else:
            default_image = get_gray_default_image(num_channels, height, width)
            metadata[name]["reshape"] = (num_channels, height, width)

        import pandas as pd

        sample_entry = abs_path_column.iloc[0] if hasattr(abs_path_column, "iloc") else next(iter(abs_path_column))
        if preprocessing_parameters.get("mode", "lazy") != "eager" and not torchvision_parameters:
            # Lazy path: store file paths as a string Series.  The actual image
            # decode happens per-batch inside PandasDataset via LazyColumn.
            # This bounds peak memory to batch_size × image_size instead of N × image_size.
            if isinstance(sample_entry, str):
                # Input is already a local/remote path — use abs_path_column directly so
                # that the DataFrame index is preserved after sampling/filtering operations.
                proc_df[feature_config[PROC_COLUMN]] = abs_path_column
            else:
                # In-memory data (HF PIL Images, raw bytes, numpy arrays) — cache to disk first.
                cache_dir = resolve_lazy_cache_dir(
                    preprocessing_parameters.get("lazy_cache_dir"),
                    name,
                )
                logger.info(f"Image feature '{name}': caching in-memory images to {cache_dir} for lazy decoding.")
                raw_column = abs_path_column.tolist() if hasattr(abs_path_column, "tolist") else list(abs_path_column)
                path_list = _cache_image_column_to_disk(raw_column, cache_dir, name)
                # Reconstruct a Series using the original index so that it aligns
                # correctly with proc_df (which may have a non-0-based index after sampling).
                if hasattr(abs_path_column, "compute"):  # Dask Series
                    orig_index = abs_path_column.index.compute()
                else:
                    orig_index = abs_path_column.index
                proc_df[feature_config[PROC_COLUMN]] = backend.df_engine.from_pandas(
                    pd.Series(path_list, dtype=object, index=orig_index)
                )
            metadata[name]["lazy"] = True  # backward compat for ray.py
            metadata[name]["mode"] = preprocessing_parameters.get("mode", "lazy")
            metadata[name]["prefetch_size"] = preprocessing_parameters.get("prefetch_size")
            metadata[name]["reshape"] = None  # paths are 1-D strings — no reshape needed
            # Persist decode params so PandasDataset can reconstruct the decode fn
            metadata[name]["lazy_image_params"] = {
                "img_width": width,
                "img_height": height,
                "should_resize": should_resize,
                "num_channels": num_channels,
                "resize_method": preprocessing_parameters["resize_method"],
                "user_specified_num_channels": user_specified_num_channels,
                "standardize_image": standardize_image,
                "channel_class_map": channel_class_map.tolist(),
                "default_image_shape": list(default_image.shape),
            }
        else:
            # Eager path (legacy): decode all images upfront into numpy arrays.
            proc_col = backend.read_binary_files(
                abs_path_column, map_fn=read_image_if_bytes_obj_and_resize, file_size=average_file_size
            )

            num_failed_image_reads = (
                proc_col.isna().sum().compute() if is_dask_series_or_df(proc_col, backend) else proc_col.isna().sum()
            )

            proc_col = backend.df_engine.map_objects(
                proc_col, lambda row: default_image if not isinstance(row, np.ndarray) else row
            )

            if num_failed_image_reads > 0:
                logger.warning(
                    f"Failed to read {num_failed_image_reads} images while preprocessing feature `{name}`. "
                    "Using default image for these rows in the dataset."
                )

            proc_df[feature_config[PROC_COLUMN]] = proc_col

        return proc_df


class ImageInputFeature(ImageFeatureMixin, InputFeature):
    def __init__(self, input_feature_config: ImageInputFeatureConfig, encoder_obj=None, **kwargs):
        super().__init__(input_feature_config, **kwargs)

        if encoder_obj:
            self.encoder_obj = encoder_obj
        else:
            self.encoder_obj = self.initialize_encoder(input_feature_config.encoder)

        # set up for augmentation if it is enabled
        if input_feature_config.augmentation:
            # assume no image normalize is required
            normalize_mean = normalize_std = None

            # determine if specified encoder is a torchvision model
            if is_torchvision_encoder(self.encoder_obj):
                # encoder is a torchvision model
                normalize_mean = self.encoder_obj.normalize_mean
                normalize_std = self.encoder_obj.normalize_std
            else:
                # encoder is a Ludwig encoder, determine if standardize_image is set to IMAGENET1K
                if input_feature_config.preprocessing.standardize_image == IMAGENET1K:
                    normalize_mean = IMAGENET1K_MEAN
                    normalize_std = IMAGENET1K_STD

            # create augmentation pipeline object
            self.augmentation_pipeline = ImageAugmentation(
                input_feature_config.augmentation,
                normalize_mean,
                normalize_std,
            )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        if not isinstance(inputs, torch.Tensor):
            raise TypeError(f"Image feature forward expects a torch.Tensor, got {type(inputs).__name__}.")
        if inputs.dtype != torch.float32:
            raise ValueError(f"Image feature inputs must be a float32 tensor, got {inputs.dtype}.")

        inputs_encoded = self.encoder_obj(inputs)

        return inputs_encoded

    @property
    def input_dtype(self):
        return torch.float32

    @property
    def input_shape(self) -> torch.Size:
        return torch.Size(self.encoder_obj.input_shape)

    @property
    def output_shape(self) -> torch.Size:
        return self.encoder_obj.output_shape

    def update_config_after_module_init(self, feature_config):
        if is_torchvision_encoder(self.encoder_obj):
            # update feature preprocessing parameters to reflect used in torchvision pretrained model
            # Note: image height and width is determined by the encoder crop_size attribute.  Source of this
            # attribute is from the torchvision.transforms._presets.ImageClassification class.  This class stores
            # crop_size as a single element list.  the single element in this list is used to set both the height
            # and width of an image.
            feature_config.preprocessing.height = self.encoder_obj.crop_size[0]
            feature_config.preprocessing.width = self.encoder_obj.crop_size[0]
            feature_config.preprocessing.num_channels = self.encoder_obj.num_channels

    @staticmethod
    def update_config_with_metadata(feature_config, feature_metadata, *args, **kwargs):
        for key in ["height", "width", "num_channels", "standardize_image"]:
            if hasattr(feature_config.encoder, key):
                setattr(feature_config.encoder, key, feature_metadata[PREPROCESSING][key])

    @staticmethod
    def get_schema_cls():
        return ImageInputFeatureConfig

    @staticmethod
    def create_preproc_module(metadata: dict[str, Any]) -> BasePreprocessingModule:
        model_type = metadata["preprocessing"].get("torchvision_model_type")
        model_variant = metadata["preprocessing"].get("torchvision_model_variant")
        if model_variant:
            torchvision_parameters = _get_torchvision_parameters(model_type, model_variant)
        else:
            torchvision_parameters = None

        if torchvision_parameters:
            torchvision_transform, transform_metadata = _get_torchvision_transform(torchvision_parameters)
        else:
            torchvision_transform = None
            transform_metadata = None

        return _ImagePreprocessing(
            metadata, torchvision_transform=torchvision_transform, transform_metadata=transform_metadata
        )

    def get_augmentation_pipeline(self):
        return self.augmentation_pipeline


class ImageOutputFeature(ImageFeatureMixin, OutputFeature):
    def __init__(
        self,
        output_feature_config: ImageOutputFeatureConfig | dict,
        output_features: dict[str, OutputFeature],
        **kwargs,
    ):
        super().__init__(output_feature_config, output_features, **kwargs)
        self.decoder_obj = self.initialize_decoder(output_feature_config.decoder)
        self._setup_loss()
        self._setup_metrics()

    def logits(self, inputs: dict[str, torch.Tensor], target=None, **kwargs):
        return self.decoder_obj(inputs, target=target)

    def metric_kwargs(self):
        return {"num_outputs": self.output_shape[0]}

    def create_predict_module(self) -> PredictModule:
        return _ImagePredict()

    def get_prediction_set(self):
        return self.decoder_obj.get_prediction_set()

    @classmethod
    def get_output_dtype(cls):
        return torch.float32

    @property
    def output_shape(self) -> torch.Size:
        return self.decoder_obj.output_shape

    @property
    def input_shape(self) -> torch.Size:
        return self.decoder_obj.input_shape

    @staticmethod
    def update_config_with_metadata(feature_config, feature_metadata, *args, **kwargs):
        for key in ["height", "width", "num_channels", "num_classes", "standardize_image"]:
            if hasattr(feature_config.decoder, key):
                setattr(feature_config.decoder, key, feature_metadata[PREPROCESSING][key])

    @staticmethod
    def calculate_overall_stats(predictions, targets, metadata):
        # no overall stats, just return empty dictionary
        return {}

    def postprocess_predictions(
        self,
        result,
        metadata,
    ):
        predictions_col = f"{self.feature_name}_{PREDICTIONS}"

        if predictions_col in result:
            channel_class_map = torch.ByteTensor(metadata[PREPROCESSING]["channel_class_map"])

            if channel_class_map.shape[0]:

                def class_mask2img(row):
                    pred = row[predictions_col]
                    return get_image_from_class_mask(channel_class_map, pred)

                result[predictions_col] = result.apply(class_mask2img, axis=1)

            return result

    @staticmethod
    def create_postproc_module(metadata: TrainingSetMetadataDict) -> torch.nn.Module:
        return _ImagePostprocessing(metadata)

    @staticmethod
    def get_schema_cls():
        return ImageOutputFeatureConfig