Coverage for python/lsst/images/_color_image.py: 85%

1# This file is part of lsst-images. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14__all__ = ("ColorImage",) 

15 

16import functools 

17from collections.abc import Sequence 

18from types import EllipsisType 

19from typing import Any, ClassVar, Literal 

20 

21import numpy as np 

22import pydantic 

23 

24from ._generalized_image import GeneralizedImage 

25from ._geom import Box 

26from ._image import Image, ImageSerializationModel 

27from ._transforms import SkyProjection, SkyProjectionSerializationModel 

28from .serialization import ArchiveTree, InputArchive, InvalidParameterError, MetadataValue, OutputArchive 

29from .utils import is_none 

30 

31 

32class ColorImage(GeneralizedImage): 

33 """An RGB image with an optional `SkyProjection`. 

34 

35 Parameters 

36 ---------- 

37 array 

38 Array or fill value for the image. Must have three dimensions with 

39 the shape of the third dimension equal to three. 

40 bbox 

41 Bounding box for the image. 

42 yx0 

43 Logical coordinates of the first pixel in the array, ordered ``y``, 

44 ``x`` (unless an `XY` instance is passed). Ignored if 

45 ``bbox`` is provided. Defaults to zeros. 

46 sky_projection 

47 Projection that maps the pixel grid to the sky. 

48 metadata 

49 Arbitrary flexible metadata to associate with the image. 

50 """ 

51 

52 def __init__( 

53 self, 

54 array: np.ndarray[tuple[int, int, Literal[3]], np.dtype[Any]], 

55 /, 

56 *, 

57 bbox: Box | None = None, 

58 yx0: Sequence[int] | None = None, 

59 sky_projection: SkyProjection[Any] | None = None, 

60 metadata: dict[str, MetadataValue] | None = None, 

61 ) -> None: 

62 super().__init__(metadata) 

63 if bbox is None: 

64 bbox = Box.from_shape(array.shape[:2], start=yx0) 

65 elif bbox.shape + (3,) != array.shape: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true

66 raise ValueError( 

67 f"Shape from bbox {bbox.shape + (3,)} does not match array with shape {array.shape}." 

68 ) 

69 self._array = array 

70 self._red = Image(self._array[..., 0], bbox=bbox, sky_projection=sky_projection) 

71 self._green = Image(self._array[..., 1], bbox=bbox, sky_projection=sky_projection) 

72 self._blue = Image(self._array[..., 2], bbox=bbox, sky_projection=sky_projection) 

73 

74 @staticmethod 

75 def from_channels( 

76 r: Image, 

77 g: Image, 

78 b: Image, 

79 *, 

80 sky_projection: SkyProjection[Any] | None = None, 

81 metadata: dict[str, MetadataValue] | None = None, 

82 ) -> ColorImage: 

83 """Construct from separate RGB images. 

84 

85 All channels are assumed to have the same bounding box, sky_projection, 

86 and pixel type. 

87 

88 Parameters 

89 ---------- 

90 r 

91 Red channel image. 

92 g 

93 Green channel image. 

94 b 

95 Blue channel image. 

96 sky_projection 

97 Sky projection for the image, defaulting to that of ``r``. 

98 metadata 

99 Flexible metadata to associate with the image. 

100 """ 

101 if sky_projection is None and r.sky_projection is not None: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true

102 sky_projection = r.sky_projection 

103 return ColorImage( 

104 np.stack([r.array, g.array, b.array], axis=2), 

105 bbox=r.bbox, 

106 sky_projection=sky_projection, 

107 metadata=metadata, 

108 ) 

109 

110 @property 

111 def array(self) -> np.ndarray[tuple[int, int, Literal[3]], np.dtype[Any]]: 

112 """The 3-d array (`numpy.ndarray`).""" 

113 return self._array 

114 

115 @property 

116 def red(self) -> Image: 

117 """A 2-d view of the red channel (`Image`).""" 

118 return self._red 

119 

120 @property 

121 def green(self) -> Image: 

122 """A 2-d view of the green channel (`Image`).""" 

123 return self._green 

124 

125 @property 

126 def blue(self) -> Image: 

127 """A 2-d view of the blue channel (`Image`).""" 

128 return self._blue 

129 

130 @property 

131 def bbox(self) -> Box: 

132 """The 2-d bounding box of the image (`Box`).""" 

133 return self._red.bbox 

134 

135 @property 

136 def sky_projection(self) -> SkyProjection[Any] | None: 

137 """The projection that maps the pixel grid to the sky 

138 (`SkyProjection` | `None`). 

139 """ 

140 return self._red.sky_projection 

141 

142 def __getitem__(self, bbox: Box | EllipsisType) -> ColorImage: 

143 bbox, indices = self._handle_getitem_args(bbox) 

144 return self._transfer_metadata( 

145 ColorImage( 

146 self.array[indices + (slice(None),)], 

147 bbox=bbox, 

148 sky_projection=self.sky_projection, 

149 ), 

150 bbox=bbox, 

151 ) 

152 

153 def __setitem__(self, bbox: Box | EllipsisType, value: ColorImage) -> None: 

154 self[bbox].array[...] = value.array 

155 

156 def __str__(self) -> str: 

157 return f"ColorImage({self.bbox!s}, {self._array.dtype.type.__name__})" 

158 

159 def __repr__(self) -> str: 

160 return f"ColorImage(..., bbox={self.bbox!r}, dtype={self._array.dtype!r})" 

161 

162 def copy(self) -> ColorImage: 

163 """Deep-copy the image.""" 

164 return self._transfer_metadata( 

165 ColorImage(self._array.copy(), bbox=self.bbox, sky_projection=self.sky_projection), copy=True 

166 ) 

167 

168 def serialize(self, archive: OutputArchive[Any]) -> ColorImageSerializationModel: 

169 """Serialize the masked image to an output archive. 

170 

171 Parameters 

172 ---------- 

173 archive 

174 Archive to write to. 

175 """ 

176 r = archive.serialize_direct("red", functools.partial(self.red.serialize, save_projection=False)) 

177 g = archive.serialize_direct("green", functools.partial(self.green.serialize, save_projection=False)) 

178 b = archive.serialize_direct("blue", functools.partial(self.blue.serialize, save_projection=False)) 

179 serialized_projection = ( 

180 archive.serialize_direct("sky_projection", self.sky_projection.serialize) 

181 if self.sky_projection is not None 

182 else None 

183 ) 

184 return ColorImageSerializationModel( 

185 red=r, green=g, blue=b, sky_projection=serialized_projection, metadata=self.metadata 

186 ) 

187 

188 @staticmethod 

189 def _get_archive_tree_type[P: pydantic.BaseModel]( 

190 pointer_type: type[P], 

191 ) -> type[ColorImageSerializationModel[P]]: 

192 """Return the serialization model type for this object for an archive 

193 type that uses the given pointer type. 

194 """ 

195 return ColorImageSerializationModel[pointer_type] # type: ignore 

196 

197 

198class ColorImageSerializationModel[P: pydantic.BaseModel](ArchiveTree): 

199 """A Pydantic model used to represent a serialized `ColorImage`.""" 

200 

201 SCHEMA_NAME: ClassVar[str] = "color_image" 

202 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

203 MIN_READ_VERSION: ClassVar[int] = 1 

204 PUBLIC_TYPE: ClassVar[type] = ColorImage 

205 

206 red: ImageSerializationModel[P] = pydantic.Field(description="The red channel.") 

207 green: ImageSerializationModel[P] = pydantic.Field(description="The green channel.") 

208 blue: ImageSerializationModel[P] = pydantic.Field(description="The blue channel") 

209 sky_projection: SkyProjectionSerializationModel[P] | None = pydantic.Field( 

210 default=None, 

211 exclude_if=is_none, 

212 description="Projection that maps the pixel grid to the sky.", 

213 ) 

214 

215 @property 

216 def bbox(self) -> Box: 

217 """The bounding box of the image.""" 

218 return self.red.bbox 

219 

220 def deserialize( 

221 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any 

222 ) -> ColorImage: 

223 """Deserialize a image from an input archive. 

224 

225 Parameters 

226 ---------- 

227 archive 

228 Archive to read from. 

229 bbox 

230 Bounding box of a subimage to read instead. 

231 **kwargs 

232 Unsupported keyword arguments are accepted only to provide 

233 better error messages (raising 

234 `.serialization.InvalidParameterError`). 

235 """ 

236 if kwargs: 236 ↛ 237line 236 didn't jump to line 237 because the condition on line 236 was never true

237 raise InvalidParameterError(f"Unrecognized parameters for ColoImage: {set(kwargs.keys())}.") 

238 r = self.red.deserialize(archive, bbox=bbox) 

239 g = self.green.deserialize(archive, bbox=bbox) 

240 b = self.blue.deserialize(archive, bbox=bbox) 

241 sky_projection = self.sky_projection.deserialize(archive) if self.sky_projection is not None else None 

242 return ColorImage.from_channels(r, g, b, sky_projection=sky_projection)._finish_deserialize(self)