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

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, 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 deserialize( 

179 model: ColorImageSerializationModel[Any], archive: InputArchive[Any], *, bbox: Box | None = None 

180 ) -> ColorImage: 

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

182 

183 Parameters 

184 ---------- 

185 model 

186 A Pydantic model representation of the image, holding references 

187 to data stored in the archive. 

188 archive 

189 Archive to read from. 

190 bbox 

191 Bounding box of a subimage to read instead. 

192 """ 

193 r = Image.deserialize(model.red, archive, bbox=bbox) 

194 g = Image.deserialize(model.green, archive, bbox=bbox) 

195 b = Image.deserialize(model.blue, archive, bbox=bbox) 

196 projection = ( 

197 Projection.deserialize(model.projection, archive) if model.projection is not None else None 

198 ) 

199 return ColorImage.from_channels(r, g, b, projection=projection)._finish_deserialize(model) 

200 

201 @staticmethod 

202 def _get_archive_tree_type[P: pydantic.BaseModel]( 

203 pointer_type: type[P], 

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

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

206 type that uses the given pointer type. 

207 """ 

208 return ColorImageSerializationModel[pointer_type] # type: ignore 

209 

210 

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

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

213 

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

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

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

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

218 default=None, 

219 exclude_if=is_none, 

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

221 ) 

222 

223 @property 

224 def bbox(self) -> Box: 

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

226 return self.red.bbox