Skip to content

dqm_ml_core.utils.image_loading

Image loading mixin for processors that load images from storage.

This module provides the ImageLoadingMixin class that handles loading images from local filesystem paths or S3, resolving path prefixes set by the job for each data selection.

logger = logging.getLogger(__name__) module-attribute

ImageLoadingMixin

Mixin for processors that load images from S3 or local paths.

Expects the host class to provide: - self.s3_fs: An S3 filesystem instance (or None). - self.current_path_prefix: A dict mapping column names to path prefixes, set by the job per data selection.

Source code in packages/dqm-ml-core/src/dqm_ml_core/utils/image_loading.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class ImageLoadingMixin:
    """Mixin for processors that load images from S3 or local paths.

    Expects the host class to provide:
    - ``self.s3_fs``: An S3 filesystem instance (or ``None``).
    - ``self.current_path_prefix``: A dict mapping column names to path
      prefixes, set by the job per data selection.
    """

    s3_fs: Any = None

    def _current_image_prefix(self, column: str | None = None) -> str | None:
        """Return the path prefix for the given column.

        Reads from ``self.current_path_prefix``, a dict set by the job
        mapping column names to path prefixes.

        Args:
            column: The input column name. Falls back to the first input
                column if ``None``.

        Returns:
            The path prefix string, or ``None``.
        """
        prefix_map: dict[str, str] = getattr(self, "current_path_prefix", {})
        if column is not None:
            return prefix_map.get(column)
        cols: list[str] = getattr(self, "input_columns", [])
        return prefix_map.get(cols[0]) if cols else None

    def _open_s3_image(self, prefix: str, path: str) -> Image.Image:
        """Open a PIL Image from S3, reading into memory and copying.

        Args:
            prefix: The path prefix within the bucket.
            path: The relative image path.

        Returns:
            A PIL Image object.
        """
        bucket = os.getenv("S3_BUCKET_NAME", "")
        s3_path = f"{bucket}/{prefix}/{path}"
        with self.s3_fs.open_input_stream(s3_path) as f:
            img = Image.open(io.BytesIO(f.read())).convert("RGB")
            return img.copy()

    def _open_image_from_path(self, path: str, column: str | None = None) -> Image.Image | None:
        """Open a PIL Image from a string path (S3 or local filesystem).

        Resolves relative paths by looking up the prefix for the configured
        input column from ``self.current_path_prefix`` (set by the job for
        each selection).

        Args:
            path: File path or relative path.
            column: The input column name (used to resolve the path prefix).

        Returns:
            A PIL Image object, or ``None`` if the file does not exist.
        """
        prefix = self._current_image_prefix(column)

        if self.s3_fs:
            return self._open_s3_image(prefix or "", path)

        img_path = Path(prefix) / path if prefix else Path(path)
        if not img_path.is_file():
            logger.warning(f"Path does not exist: {img_path}")
            return None
        return Image.open(img_path)

s3_fs: Any = None class-attribute instance-attribute