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

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 super().__getitem__(bbox) 

144 if bbox is ...: 

145 return self 

146 return self._transfer_metadata( 

147 ColorImage( 

148 self.array[bbox.slice_within(self.bbox) + (slice(None),)], 

149 bbox=bbox, 

150 sky_projection=self.sky_projection, 

151 ), 

152 bbox=bbox, 

153 ) 

154 

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

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

157 

158 def __str__(self) -> str: 

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

160 

161 def __repr__(self) -> str: 

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

163 

164 def copy(self) -> ColorImage: 

165 """Deep-copy the image.""" 

166 return self._transfer_metadata( 

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

168 ) 

169 

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

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

172 

173 Parameters 

174 ---------- 

175 archive 

176 Archive to write to. 

177 """ 

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

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

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

181 serialized_projection = ( 

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

183 if self.sky_projection is not None 

184 else None 

185 ) 

186 return ColorImageSerializationModel( 

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

188 ) 

189 

190 @staticmethod 

191 def _get_archive_tree_type[P: pydantic.BaseModel]( 

192 pointer_type: type[P], 

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

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

195 type that uses the given pointer type. 

196 """ 

197 return ColorImageSerializationModel[pointer_type] # type: ignore 

198 

199 

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

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

202 

203 SCHEMA_NAME: ClassVar[str] = "color_image" 

204 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

205 MIN_READ_VERSION: ClassVar[int] = 1 

206 PUBLIC_TYPE: ClassVar[type] = ColorImage 

207 

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

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

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

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

212 default=None, 

213 exclude_if=is_none, 

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

215 ) 

216 

217 @property 

218 def bbox(self) -> Box: 

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

220 return self.red.bbox 

221 

222 def deserialize( 

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

224 ) -> ColorImage: 

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

226 

227 Parameters 

228 ---------- 

229 archive 

230 Archive to read from. 

231 bbox 

232 Bounding box of a subimage to read instead. 

233 **kwargs 

234 Unsupported keyword arguments are accepted only to provide 

235 better error messages (raising 

236 `.serialization.InvalidParameterError`). 

237 """ 

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

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

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

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

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

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

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