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

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, 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 Projection, ProjectionSerializationModel 

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 `Projection`. 

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 start 

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 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 start: Sequence[int] | None = None, 

59 projection: Projection[Any] | None = None, 

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

61 ): 

62 super().__init__(metadata) 

63 if bbox is None: 

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

65 elif bbox.shape + (3,) != array.shape: 

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, projection=projection) 

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

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

73 

74 @staticmethod 

75 def from_channels( 

76 r: Image, 

77 g: Image, 

78 b: Image, 

79 *, 

80 projection: Projection[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, projection, 

86 and pixel type. 

87 """ 

88 if projection is None and r.projection is not None: 

89 projection = r.projection 

90 return ColorImage( 

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

92 bbox=r.bbox, 

93 projection=projection, 

94 metadata=metadata, 

95 ) 

96 

97 @property 

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

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

100 return self._array 

101 

102 @property 

103 def red(self) -> Image: 

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

105 return self._red 

106 

107 @property 

108 def green(self) -> Image: 

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

110 return self._green 

111 

112 @property 

113 def blue(self) -> Image: 

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

115 return self._blue 

116 

117 @property 

118 def bbox(self) -> Box: 

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

120 return self._red.bbox 

121 

122 @property 

123 def projection(self) -> Projection[Any] | None: 

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

125 (`Projection` | `None`). 

126 """ 

127 return self._red.projection 

128 

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

130 super().__getitem__(bbox) 

131 if bbox is ...: 

132 return self 

133 return self._transfer_metadata( 

134 ColorImage( 

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

136 bbox=bbox, 

137 projection=self.projection, 

138 ), 

139 bbox=bbox, 

140 ) 

141 

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

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

144 

145 def __str__(self) -> str: 

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

147 

148 def __repr__(self) -> str: 

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

150 

151 def copy(self) -> ColorImage: 

152 """Deep-copy the image.""" 

153 return self._transfer_metadata( 

154 ColorImage(self._array.copy(), bbox=self.bbox, projection=self.projection), copy=True 

155 ) 

156 

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

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

159 

160 Parameters 

161 ---------- 

162 archive 

163 Archive to write to. 

164 """ 

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

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

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

168 serialized_projection = ( 

169 archive.serialize_direct("projection", self.projection.serialize) 

170 if self.projection is not None 

171 else None 

172 ) 

173 return ColorImageSerializationModel( 

174 red=r, green=g, blue=b, projection=serialized_projection, metadata=self.metadata 

175 ) 

176 

177 @staticmethod 

178 def _get_archive_tree_type[P: pydantic.BaseModel]( 

179 pointer_type: type[P], 

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

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

182 type that uses the given pointer type. 

183 """ 

184 return ColorImageSerializationModel[pointer_type] # type: ignore 

185 

186 

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

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

189 

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

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

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

193 projection: ProjectionSerializationModel[P] | None = pydantic.Field( 

194 default=None, 

195 exclude_if=is_none, 

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

197 ) 

198 

199 @property 

200 def bbox(self) -> Box: 

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

202 return self.red.bbox 

203 

204 def deserialize( 

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

206 ) -> ColorImage: 

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

208 

209 Parameters 

210 ---------- 

211 model 

212 A Pydantic model representation of the image, holding references 

213 to data stored in the archive. 

214 archive 

215 Archive to read from. 

216 bbox 

217 Bounding box of a subimage to read instead. 

218 """ 

219 if kwargs: 

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

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

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

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

224 projection = self.projection.deserialize(archive) if self.projection is not None else None 

225 return ColorImage.from_channels(r, g, b, projection=projection)._finish_deserialize(self)