Coverage for packages/dqm-ml-core/src/dqm_ml_core/utils/image_loading.py: 81%

1"""Image loading mixin for processors that load images from storage. 

2 

3This module provides the ImageLoadingMixin class that handles loading 

4images from local filesystem paths or S3, resolving path prefixes 

5set by the job for each data selection. 

6""" 

7 

8from __future__ import annotations 

9 

10import io 

11import logging 

12import os 

13from pathlib import Path 

14from typing import Any 

15 

16from PIL import Image 

17 

18logger = logging.getLogger(__name__) 

19 

20 

21class ImageLoadingMixin: 

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

23 

24 Expects the host class to provide: 

25 - ``self.s3_fs``: An S3 filesystem instance (or ``None``). 

26 - ``self.current_path_prefix``: A dict mapping column names to path 

27 prefixes, set by the job per data selection. 

28 """ 

29 

30 s3_fs: Any = None 

31 

32 def _current_image_prefix(self, column: str | None = None) -> str | None: 

33 """Return the path prefix for the given column. 

34 

35 Reads from ``self.current_path_prefix``, a dict set by the job 

36 mapping column names to path prefixes. 

37 

38 Args: 

39 column: The input column name. Falls back to the first input 

40 column if ``None``. 

41 

42 Returns: 

43 The path prefix string, or ``None``. 

44 """ 

45 prefix_map: dict[str, str] = getattr(self, "current_path_prefix", {}) 

46 if column is not None: 

47 return prefix_map.get(column) 

48 cols: list[str] = getattr(self, "input_columns", []) 

49 return prefix_map.get(cols[0]) if cols else None 

50 

51 def _open_s3_image(self, prefix: str, path: str) -> Image.Image: 

52 """Open a PIL Image from S3, reading into memory and copying. 

53 

54 Args: 

55 prefix: The path prefix within the bucket. 

56 path: The relative image path. 

57 

58 Returns: 

59 A PIL Image object. 

60 """ 

61 bucket = os.getenv("S3_BUCKET_NAME", "") 

62 s3_path = f"{bucket}/{prefix}/{path}" 

63 with self.s3_fs.open_input_stream(s3_path) as f: 

64 img = Image.open(io.BytesIO(f.read())).convert("RGB") 

65 return img.copy() 

66 

67 def _open_image_from_path(self, path: str, column: str | None = None) -> Image.Image | None: 

68 """Open a PIL Image from a string path (S3 or local filesystem). 

69 

70 Resolves relative paths by looking up the prefix for the configured 

71 input column from ``self.current_path_prefix`` (set by the job for 

72 each selection). 

73 

74 Args: 

75 path: File path or relative path. 

76 column: The input column name (used to resolve the path prefix). 

77 

78 Returns: 

79 A PIL Image object, or ``None`` if the file does not exist. 

80 """ 

81 prefix = self._current_image_prefix(column) 

82 

83 if self.s3_fs: 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true

84 return self._open_s3_image(prefix or "", path) 

85 

86 img_path = Path(prefix) / path if prefix else Path(path) 

87 if not img_path.is_file(): 

88 logger.warning(f"Path does not exist: {img_path}") 

89 return None 

90 return Image.open(img_path)