Coverage for collectra/types/images.py: 77.46%

1"""Image handling and processing classes for the Collectra system.

3This module provides the core image classes used throughout Collectra for

4handling image data, metadata, and cropping operations. It supports various

5image formats and provides utilities for encoding, metadata extraction,

6and coordinate-based cropping.

8The module includes:

9 - Base Image class for general image handling

10 - ImageCrop class for coordinate-based image cropping

11 - Image format validation and metadata extraction

12 - Base64 encoding for image serialization

14Classes:

15 Image: Base class for image handling and metadata

16 ImageCrop: Specialized class for cropped image regions

17"""

19__all__ = ["Image", "ImageCrop", "ImageSegmentation"]

22import base64

23import io

24import shutil

25from dataclasses import dataclass, field

26from enum import Enum

27from functools import lru_cache

28from pathlib import Path

30from PIL import Image as ImagePil

32from collectra.cli import command

34from .base import Artefact, ArtefactCommands, ArtefactNode

37def _tiff_dimensions(path: str | Path) -> tuple[int, int, int]:

38 """Return width, height, channels from TIFF headers, without decoding pixels.

40 C/S denote channels/samples; I/Q denotes an otherwise unlabelled page stack.

41 Non-singleton time/depth axes and multiple image series are ambiguous for a

42 2D detector and must be exported as separate images first.

43 """

44 from tifffile import TiffFile

46 with TiffFile(path) as tiff:

47 if len(tiff.series) != 1:

48 raise ValueError(f"TIFF {path} contains multiple image series")

49 series = tiff.series[0]

50 shape, axes = series.shape, series.axes

51 if len(axes) != len(shape) or axes.count("Y") != 1 or axes.count("X") != 1:

52 raise ValueError(

53 f"Cannot identify TIFF spatial axes for {path}: shape={shape}, axes={axes!r}"

54 )

55 extra = [

56 (axis, size)

57 for axis, size in zip(axes, shape)

58 if axis not in "YX" and size != 1

59 ]

60 if len(extra) > 1 or (extra and extra[0][0] not in "CSIQ"):

61 raise ValueError(

62 f"Ambiguous TIFF channel axes for {path}: shape={shape}, axes={axes!r}; expected a 2D image with one channel/sample axis"

63 )

64 channels = extra[0][1] if extra else 1

65 return shape[axes.index("X")], shape[axes.index("Y")], channels

68def image_channel_count(path: str | Path) -> int:

69 """Inspect image metadata for its channel count; never convert image pixels."""

70 if Path(path).suffix.lower() in {".tif", ".tiff"}:

71 return _tiff_dimensions(path)[2]

72 with ImagePil.open(path) as image:

73 if image.mode == "P":

74 return 4 if "transparency" in image.info else 3

75 return len(image.getbands())

78def image_size(path: str | Path) -> tuple[int, int]:

79 """Return ``(width, height)`` from image metadata; supports multi-channel TIFFs."""

80 if Path(path).suffix.lower() in {".tif", ".tiff"}:

81 return _tiff_dimensions(path)[:2]

82 with ImagePil.open(path) as image:

83 return image.size

86def read_tiff_channels(path: str | Path, maxworkers: int | None = None):

87 """Decode a TIFF into HWC order using the same axis rules as metadata inspection."""

88 from tifffile import TiffFile

90 width, height, channels = _tiff_dimensions(path)

91 with TiffFile(path) as tiff:

92 series = tiff.series[0]

93 axes = series.axes

94 if maxworkers is None:

95 pixels = series.asarray()

96 else:

97 try:

98 pixels = series.asarray(maxworkers=maxworkers)

99 except TypeError:

100 # Older tifffile versions may not expose maxworkers.

101 pixels = series.asarray()

102 spatial = [axes.index("Y"), axes.index("X")]

103 order = spatial + [index for index in range(pixels.ndim) if index not in spatial]

104 return pixels.transpose(order).reshape(height, width, channels)

105

106

107@lru_cache(maxsize=4096)

108def _image_metadata(

109 path: str, modified_ns: int, file_size: int

110) -> tuple[int, int, str | None]:

111 """Read image dimensions and format once per file version."""

112 if Path(path).suffix.lower() in {".tif", ".tiff"}:

113 width, height, _ = _tiff_dimensions(path)

114 return width, height, "TIFF"

115 with ImagePil.open(path) as image:

116 width, height = image.size

117 return width, height, image.format

118

119

120class ImageArtefactCommands(ArtefactCommands):

121 """Commands applying to image artefacts stored in a pipeline node."""

122

123 @command

124 def extract(

125 self,

126 *inputs: str,

127 output: Path = Path("extracted"),

128 ) -> None:

129 """Extract this artefact from Collectra result folders."""

130 output = output.expanduser().resolve()

131 output.mkdir(parents=True, exist_ok=True)

132

133 sources: list[tuple[str, dict[str, Artefact]]] = []

134 if self.node.items:

135 sources.append(("configured", self.node.items))

136

137 for input_path in inputs:

138 path = Path(input_path).expanduser().resolve()

139 loaded = ArtefactNode(self.node.name, types=set(self.node.types))

140 loaded.process(self.node.name, path)

141 sources.append((path.stem, loaded.items))

142

143 for source_name, items in sources:

144 destination = output / source_name if len(sources) > 1 else output

145 destination.mkdir(parents=True, exist_ok=True)

146 for item in items.values():

147 if isinstance(item, Image):

148 suffix = ".png" if isinstance(item, ImageSegmentation) else ".jpg"

149 item.extract(destination / f"{item.id}{suffix}")

150

151

152class Orientation(Enum):

153 """Enumeration for image orientation states."""

154

155 NORTH = 0

156 WEST = 1

157 SOUTH = 2

158 EAST = 3

159

160 def to_string(self) -> str:

161 match self.value:

162 case 0:

163 return "north"

164 case 1:

165 return "west"

166 case 2:

167 return "south"

168 case 3:

169 return "east"

170

171 def to_degree(self) -> int:

172 match self.value:

173 case 0:

174 return 0

175 case 1:

176 return -90

177 case 2:

178 return -180

179 case 3:

180 return -270

181

182 @staticmethod

183 def from_string(direction: str) -> "Orientation":

184 direction = direction.strip().lower()

185 match direction:

186 case "north":

187 return Orientation.NORTH

188 case "west":

189 return Orientation.WEST

190 case "south":

191 return Orientation.SOUTH

192 case "east":

193 return Orientation.EAST

194 case _:

195 raise ValueError(f"Invalid degree for orientation: {direction}")

196

197

198@dataclass

199class Image(Artefact):

200 """Base class for handling image data and metadata in Collectra workflows.

201

202 Provides core functionality for loading, encoding, and managing image files

203 with support for various formats and metadata extraction. Serves as the

204 foundation for more specialized image processing classes.

205

206 Attributes:

207 path (str): File system path to the image file.

208 raw_width (int): Width of the image in pixels.

209 raw_height (int): Height of the image in pixels.

210 format (str | None): Image format (PNG, JPEG, etc.) or None if unknown.

211 """

212

213 data: str | Path | ImagePil.Image = field(default="") # Path to the image file

214 raw_width: int = field(init=False, default=0) # Image width in pixels

215 raw_height: int = field(init=False, default=0) # Image height in pixels

216 ext: str | None = field(default="") # Image format (e.g., PNG, JPEG)

217 embeddings: str | list[str] = field(default_factory=list) # Optional embedding data

218 orientation: Orientation = field(default=Orientation.NORTH) # Image orientation

219

220 @classmethod

221 def cli_commands(cls, node: ArtefactNode, ext: str) -> ImageArtefactCommands:

222 """Expose image operations for a named pipeline artefact node."""

223 return ImageArtefactCommands(node, ext)

224

225 def attributes_to_ignore(self):

226 attributes = super().attributes_to_ignore()

227 [attributes.add(attr) for attr in ["raw_width", "raw_height", "ext"]]

228 return attributes

229

230 def serialize(self) -> dict:

231 serialized = super().serialize()

232 serialized["orientation"] = (

233 self.orientation.to_string()

234 if isinstance(self.orientation, Orientation)

235 else self.orientation

236 )

237 return serialized

238

239 def save(self, path: Path | str = None):

240 path = Path(path)

241 path.parent.mkdir(exist_ok=True, parents=True)

242

243 # Simply copy to new dest

244 shutil.copy(self.get_path(), path)

245

246 @staticmethod

247 def image_types() -> list[str]:

248 return [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tif", ".tiff", ".webp"]

249

250 @property

251 def width(self) -> int | float:

252 return self.raw_width

253

254 @property

255 def height(self) -> int | float:

256 return self.raw_height

257

258 def __post_init__(self):

259 """Create an Image instance from a file path.

260

261 Loads image metadata (dimensions and format) from the specified file

262 and returns a configured Image instance.

263

264 Args:

265 data (Path): Path to the image file to load.

266

267 Returns:

268 Image: Configured Image instance with loaded metadata.

269

270 Raises:

271 FileNotFoundError: If the image file doesn't exist.

272 PIL.UnidentifiedImageError: If the file is not a valid image.

273 """

274 super().__post_init__()

275 if not self.data:

276 raise ValueError("Image data path is empty.")

277 if not isinstance(self.data, (str, Path)):

278 raise TypeError("Image data must be a file path.")

279 if not isinstance(self.data, Path):

280 self.data = Path.cwd() / self.data

281 if not self.data.exists() or not self.data.is_file():

282 raise FileNotFoundError(f"Image file not found: {self.data}")

283 if not isinstance(self.embeddings, list):

284 self.embeddings = [self.embeddings]

285 if isinstance(self.orientation, str):

286 self.orientation = Orientation.from_string(self.orientation)

287 stat = self.data.stat()

288 self.raw_width, self.raw_height, self.ext = _image_metadata(

289 str(self.data.resolve()), stat.st_mtime_ns, stat.st_size

290 )

291

292 def __call__(self) -> str:

293 return str(self.get_path())

294

295 def get_path(self) -> Path:

296 """Get the file path of the image.

297

298 Returns:

299 Path: The file path of the image.

300 """

301 if not isinstance(self.data, Path):

302 if not isinstance(self.data, Image):

303 raise Exception(

304 "This appear to be not an Image object. Artefact path is only available for Image object."

305 )

306 raise Exception("Image data is not a valid Path object.")

307 return self.data

308

309 @staticmethod

310 def is_image_file(path: Path) -> bool:

311 """Check if a file path represents a supported image format.

312

313 Args:

314 path (Path): File path to check.

315

316 Returns:

317 bool: True if the file extension indicates a supported image format.

318 """

319 return path.suffix.lower() in Image.image_types()

320

321 def _load_buffer(self, img: ImagePil.Image) -> bytes:

322 buffer = io.BytesIO()

323 img.save(buffer, format=self.ext)

324 buffer.seek(0)

325 return buffer.read()

326

327 def load(self) -> bytes:

328 """Load the image file content as raw bytes.

329

330 Returns:

331 bytes: Raw binary content of the image file.

332

333 Raises:

334 FileNotFoundError: If the image file doesn't exist.

335 PermissionError: If the file cannot be read due to permissions.

336 """

337 img = self.pil()

338 return self._load_buffer(img)

339

340 def get_encoding(self) -> str:

341 """Get the base64 encoded representation of the image.

342

343 Loads the image content and encodes it as a base64 string,

344 suitable for embedding in JSON or other text-based formats.

345

346 Returns:

347 str: Base64 encoded string representation of the image.

348 """

349 return base64.b64encode(self.load()).decode("utf-8")

350

351 def mime(self) -> str:

352 """Get the MIME type for the image format.

353

354 Returns:

355 str: MIME type string (e.g., 'image/jpeg', 'image/png').

356 """

357 return f"image/{self.ext.lower()}" if self.ext else "image"

358

359 def pil(self) -> ImagePil.Image:

360 image = ImagePil.open(self.get_path())

361 return image.rotate(self.orientation.to_degree(), expand=True)

362

363 def display(self, context) -> dict:

364 return context.image(self)

365

366 def display_bounds(self):

367 """Region in source pixels represented by this image's viewer."""

368 return (0, 0, self.raw_width, self.raw_height)

369

370 def _extract(self, path: Path) -> None:

371 """Extract the image to a specified path."""

372 if path.suffix == "":

373 path = path.with_suffix(".jpg")

374 self.pil().save(path)

375

376 def check_valid_relative_crop_values(

377 self,

378 x_center: float,

379 y_center: float,

380 width_relative: float,

381 height_relative: float,

382 ):

383 if not (0.0 <= x_center <= 1.0):

384 raise ValueError(f"x_center is {x_center}")

385

386 if not (0.0 <= y_center <= 1.0):

387 raise ValueError(f"y_center is {y_center}")

388

389 if not (0.0 <= width_relative <= 1.0):

390 raise ValueError(f"width_relative is {width_relative}")

391

392 if not (0.0 <= height_relative <= 1.0):

393 raise ValueError(f"height_relative is {height_relative}")

394

395 def make_crop(

396 self,

397 x_center: float,

398 y_center: float,

399 width_relative: float,

400 height_relative: float,

401 confidence: float | None = None,

402 orientation: Orientation = Orientation.NORTH,

403 name: str = "",

404 ) -> "ImageCrop":

405

406 self.check_valid_relative_crop_values(

407 x_center,

408 y_center,

409 width_relative,

410 height_relative,

411 )

412

413 return ImageCrop(

414 name=name if name else self.name,

415 data=self.data,

416 x_center=x_center,

417 y_center=y_center,

418 width_relative=width_relative,

419 height_relative=height_relative,

420 confidence=confidence,

421 orientation=orientation,

422 )

423

424 def make_crop_bounding_box(

425 self,

426 left,

427 top,

428 right,

429 bottom,

430 min_height: float = 0.0,

431 confidence: float | None = None,

432 ) -> "ImageCrop":

433 bbox_width = right - left

434 bbox_height = bottom - top

435 x_center = (left + 0.5 * bbox_width) / self.width

436 y_center = (top + 0.5 * bbox_height) / self.height

437 width_relative = bbox_width / self.width

438 height_relative = bbox_height / self.height

439

440 return self.make_crop(

441 x_center=x_center,

442 y_center=y_center,

443 width_relative=width_relative,

444 height_relative=height_relative,

445 confidence=confidence,

446 )

447

448 def evaluate(self, gold: "ImageCrop") -> float:

449 raise NotImplementedError(

450 "Base Image class does not implement evaluate(). Use ImageCrop instead."

451 )

452

453 def set_rel_to_src_parent(self) -> None:

454 pass

455

456

457@dataclass

458class ImageCrop(Image):

459 """Specialized image class for handling cropped regions of images.

460

461 Extends the base Image class to support coordinate-based cropping with

462 relative positioning. Coordinates are specified as normalized values

463 (0.0 to 1.0) relative to the original image dimensions.

464

465 Attributes:

466 x_center (float): Normalized x-coordinate of crop center (0.0 to 1.0).

467 y_center (float): Normalized y-coordinate of crop center (0.0 to 1.0).

468 width_relative (float): Normalized width of crop region (0.0 to 1.0).

469 height_relative (float): Normalized height of crop region (0.0 to 1.0).

470 """

471

472 x_center: float = field(default=0.5)

473 y_center: float = field(default=0.5)

474 width_relative: float = field(default=1.0)

475 height_relative: float = field(default=1.0)

476 confidence: float | None = field(default=None)

477 source_parent: "Image | ImageCrop | None" = field(init=False, default=None)

478

479 def display_bounds(self):

480 return self.coordinates()

481

482 def save(self, path: Path | str = None):

483 im = self.source_parent.pil() if self.source_parent else self.pil()

484 if path is None:

485 raise ValueError("Path must be provided to save the cropped image.")

486 im.save(path)

487

488 @property

489 def width(self):

490 return self.width_relative * self.raw_width

491

492 @property

493 def height(self):

494 return self.height_relative * self.raw_height

495

496 def add_source_parent(self, parent: "Image | ImageCrop") -> None:

497 self.source_parent = parent

498

499 def set_rel_to_src_parent(self) -> None:

500

501 if not self.source_parent:

502 return

503

504 if type(self.source_parent) == ImageCrop:

505 self.source_parent.set_rel_to_src_parent()

506 dx = self.x_center - self.source_parent.x_center

507 dy = self.y_center - self.source_parent.y_center

508 self.x_center = 0.5 + dx / self.source_parent.width_relative

509 self.y_center = 0.5 + dy / self.source_parent.height_relative

510 self.width_relative /= self.source_parent.width_relative

511 self.height_relative /= self.source_parent.height_relative

512

513 match self.source_parent.orientation.to_degree():

514 case -90:

515 self.x_center, self.y_center = 1 - self.y_center, self.x_center

516 self.width_relative, self.height_relative = (

517 self.height_relative,

518 self.width_relative,

519 )

520 case -180:

521 self.x_center, self.y_center = 1 - self.x_center, 1 - self.y_center

522 case -270:

523 self.x_center, self.y_center = self.y_center, 1 - self.x_center

524 self.width_relative, self.height_relative = (

525 self.height_relative,

526 self.width_relative,

527 )

528

529 def pil(self) -> ImagePil.Image:

530 coordinates = self.coordinates()

531 left, upper, right, bottom = coordinates

532 if right <= left or bottom <= upper:

533 raise ValueError(

534 f"ImageCrop {self.id!r} produces an empty pixel crop: "

535 f"input_file={str(self.source_file) if self.source_file else 'unknown'!r}, "

536 f"image_file={str(self.get_path())!r}, raw_size=({self.raw_width}, "

537 f"{self.raw_height}), normalized=(x_center={self.x_center}, "

538 f"y_center={self.y_center}, width={self.width_relative}, "

539 f"height={self.height_relative}), pixel_box={coordinates}"

540 )

541 with ImagePil.open(self.get_path()) as imf:

542 im_crop = imf.crop(coordinates)

543 return im_crop.rotate(self.orientation.to_degree(), expand=True)

544

545 def coordinates(self) -> tuple[float, float, float, float]:

546 """Calculate absolute pixel coordinates for the crop region.

547

548 Converts the normalized crop coordinates to absolute pixel coordinates

549 within the bounds of the original image dimensions.

550

551 Returns:

552 tuple[float, float, float, float]: Absolute coordinates as

553 (left, upper, right, bottom) in pixels.

554 """

555 x_center = self.x_center * self.raw_width

556 y_center = self.y_center * self.raw_height

557 actual_width_half = self.raw_width * self.width_relative / 2

558 actual_height_half = self.raw_height * self.height_relative / 2

559 left = max(0, min(int(x_center - actual_width_half), self.raw_width))

560 upper = max(0, min(int(y_center - actual_height_half), self.raw_height))

561 right = max(0, min(int(x_center + actual_width_half), self.raw_width))

562 bottom = max(0, min(int(y_center + actual_height_half), self.raw_height))

563

564 return (left, upper, right, bottom)

565

566 def make_crop(

567 self,

568 x_center: float,

569 y_center: float,

570 width_relative: float,

571 height_relative: float,

572 confidence: float | None = None,

573 orientation: Orientation = Orientation.NORTH,

574 name: str = "",

575 ) -> "ImageCrop":

576

577 degree_transformed = orientation.to_degree()

578

579 if degree_transformed == -90:

580 x_center, y_center = y_center, 1 - x_center

581 width_relative, height_relative = height_relative, width_relative

582 elif degree_transformed == -180:

583 x_center, y_center = 1 - x_center, 1 - y_center

584 elif degree_transformed == -270:

585 x_center, y_center = 1 - y_center, x_center

586 width_relative, height_relative = height_relative, width_relative

587

588 name = name if name else self.name

589

590 return super().make_crop(

591 x_center=self.x_center + (x_center - 0.5) * self.width_relative,

592 y_center=self.y_center + (y_center - 0.5) * self.height_relative,

593 width_relative=width_relative * self.width_relative,

594 height_relative=height_relative * self.height_relative,

595 confidence=confidence,

596 orientation=orientation,

597 name=name,

598 )

599

600 @staticmethod

601 def compute_iou(boxA, boxB):

602 # determine the (x, y)-coordinates of the intersection rectangle

603 xA = max(boxA[0], boxB[0])

604 yA = max(boxA[1], boxB[1])

605 xB = min(boxA[2], boxB[2])

606 yB = min(boxA[3], boxB[3])

607

608 # compute the area of intersection rectangle

609 interWidth = max(0, xB - xA)

610 interHeight = max(0, yB - yA)

611 interArea = interWidth * interHeight

612

613 # compute the area of both the prediction and ground-truth rectangles

614 boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])

615 boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])

616

617 # compute the intersection over union by taking the intersection

618 # area and dividing it by the sum of prediction + ground-truth

619 # areas - the interesection area

620 if boxAArea + boxBArea - interArea == 0:

621 return 0.0

622 iou = interArea / float(boxAArea + boxBArea - interArea)

623

624 # return the intersection over union value

625 return iou

626

627 def evaluate(self, gold: "ImageCrop") -> float:

628 if not isinstance(gold, ImageCrop):

629 raise ValueError("Reference data must be an instance of ImageCrop.")

630 iou = 0.0

631 boxA = self.coordinates()

632 boxB = gold.coordinates()

633 iou = self.compute_iou(boxA, boxB)

634 return iou

635

636

637@dataclass

638class ImageSegmentation(Image):

639 """An image region defined by a source-sized binary PNG mask.

640

641 Mask pixels are 0 outside the region and 1 inside. Masks in Pillow's

642 one-bit mode are also supported. The rendered image is cropped to the

643 nonzero bounding box and retains source alpha within the region.

644 """

645

646 mask: str | Path = field(default="")

647

648 def __post_init__(self):

649 super().__post_init__()

650 if not self.mask:

651 raise ValueError("ImageSegmentation mask path is empty")

652 self.mask = Path(self.mask)

653 self._mask_path = self.mask.resolve()

654 self._mask_image()

655

656 def attributes_to_ignore(self):

657 return super().attributes_to_ignore() | {"_mask_path"}

658

659 def _mask_image(self) -> ImagePil.Image:

660 import numpy as np

661

662 with ImagePil.open(self._mask_path) as mask:

663 if mask.format != "PNG":

664 raise ValueError("ImageSegmentation mask must be a PNG file")

665 if mask.size != (self.raw_width, self.raw_height):

666 raise ValueError("Mask dimensions must match the source image")

667 pixels = np.asarray(mask)

668 if pixels.ndim != 2 or not np.isin(pixels, [0, 1]).all():

669 raise ValueError("Mask must contain only single-channel 0 and 1 pixels")

670 alpha = ImagePil.fromarray(pixels.astype("uint8") * 255)

671 if alpha.getbbox() is None:

672 raise ValueError("ImageSegmentation mask is empty")

673 return alpha

674

675 def serialize(self) -> dict:

676 serialized = super().serialize()

677 serialized["mask"] = str(self.mask)

678 return serialized

679

680 def display_bounds(self):

681 return self._mask_image().getbbox()

682

683 @property

684 def width(self):

685 left, top, right, bottom = self.display_bounds()

686 return right - left

687

688 @property

689 def height(self):

690 left, top, right, bottom = self.display_bounds()

691 return bottom - top

692

693 def masked_pil(self, image: ImagePil.Image) -> ImagePil.Image:

694 """Mask and crop an unrotated source image or display projection."""

695 from PIL import ImageChops

696

697 mask = self._mask_image()

698 if image.size != mask.size:

699 raise ValueError("Mask dimensions must match the source image")

700 result = image.convert("RGBA")

701 result.putalpha(ImageChops.multiply(result.getchannel("A"), mask))

702 return result.crop(mask.getbbox()).rotate(

703 self.orientation.to_degree(), expand=True

704 )

705

706 def pil(self) -> ImagePil.Image:

707 with ImagePil.open(self.get_path()) as image:

708 return self.masked_pil(image)

709

710 def _load_buffer(self, img: ImagePil.Image) -> bytes:

711 buffer = io.BytesIO()

712 img.save(buffer, format="PNG")

713 return buffer.getvalue()

714

715 def mime(self) -> str:

716 return "image/png"

717

718 def save(self, path: Path | str = None):

719 if path is None:

720 raise ValueError("Path must be provided to save the segmented image")

721 self.extract(path)

722

723 def _extract(self, path: Path) -> None:

724 if not path.suffix:

725 path = path.with_suffix(".png")

726 self.pil().save(path)