Coverage for src/arcade_collection/convert/convert_to_images.py: 100%
75 statements
« prev ^ index » next coverage.py v7.1.0, created at 2024-12-09 19:07 +0000
1from __future__ import annotations
3from enum import Enum
4from typing import TYPE_CHECKING
6import numpy as np
8from arcade_collection.output.extract_tick_json import extract_tick_json
9from arcade_collection.output.get_location_voxels import get_location_voxels
11if TYPE_CHECKING:
12 import tarfile
15class ImageType(Enum):
16 """Image conversion types."""
18 FULL = (False, False, False, False)
19 """Image with TCZYX dimensions."""
21 FULL_BINARY = (True, False, False, False)
22 """Binary image with TCZYX dimensions."""
24 FULL_BY_FRAME = (False, True, False, False)
25 """Image with TCZYX dimensions separated by frame."""
27 FULL_BINARY_BY_FRAME = (True, True, False, False)
28 """Binary image with TCZYX dimensions separated by frame."""
30 FLAT_BY_FRAME = (False, True, True, False)
31 """Image array flattened to YX dimensions separated by frame."""
33 FLAT_BINARY_BY_FRAME = (True, True, True, False)
34 """Binary array flattened to YX dimensions separated by frame."""
36 FLAT_RGBA_BY_FRAME = (False, True, True, True)
37 """RGBA array flattened to YX dimensions separated by frame ."""
40def convert_to_images(
41 series_key: str,
42 locations_tar: tarfile.TarFile,
43 frame_spec: tuple[int, int, int],
44 regions: list[str],
45 box: tuple[int, int, int],
46 chunk_size: int,
47 image_type: ImageType,
48) -> list[tuple[int, int, np.ndarray, int | None]]:
49 """
50 Convert data to image arrays.
52 Images are extracted from lists of voxels. The initial converted image has
53 dimensions in TCZYX order, such that T encodes the specified frames and C
54 encodes the regions. The initial converted image is then further processed
55 based on selected image type.
57 Parameters
58 ----------
59 series_key
60 Simulation series key.
61 locations_tar
62 Archive of location data.
63 frame_spec
64 Specification for image frames.
65 regions
66 List of region channels.
67 box
68 Size of bounding box.
69 chunk_size
70 Size of each image chunk.
71 image_type
72 Image conversion type.
74 Returns
75 -------
76 :
77 List of image chunks, chunk indices, and frames.
78 """
80 binary, separate, _, reindex = image_type.value
81 length, width, height = box
82 frames = list(np.arange(*frame_spec))
83 raw_array = np.zeros((len(frames), len(regions), height, width, length), "uint16")
85 object_id = 1
87 for index, frame in enumerate(frames):
88 locations = extract_tick_json(locations_tar, series_key, frame, "LOCATIONS")
90 for location in locations:
91 value = object_id if binary or reindex else location["id"]
93 for channel, region in enumerate(regions):
94 voxels = [
95 (z, y, x)
96 for x, y, z in get_location_voxels(
97 location, region if region != "DEFAULT" else None
98 )
99 ]
101 if len(voxels) == 0:
102 continue
104 raw_array[index, channel][tuple(np.transpose(voxels))] = value
106 if reindex:
107 object_id = object_id + 1
109 # Remove 1 pixel border.
110 if height == 1:
111 array = raw_array[:, :, :, 1:-1, 1:-1].copy()
112 else:
113 array = raw_array[:, :, 1:-1, 1:-1, 1:-1].copy()
115 if separate:
116 chunks = [
117 (i, j, flatten_array_chunk(chunk, image_type), frame)
118 for index, frame in enumerate(frames)
119 for i, j, chunk in split_array_chunks(array[[index], :, :, :, :], chunk_size)
120 ]
121 else:
122 chunks = [(i, j, chunk, None) for i, j, chunk in split_array_chunks(array, chunk_size)]
124 return chunks
127def split_array_chunks(array: np.ndarray, chunk_size: int) -> list[tuple[int, int, np.ndarray]]:
128 """
129 Split arrays into smaller chunks.
131 Parameters
132 ----------
133 array
134 Image array (dimensions in TCZYX order).
135 chunk_size
136 Size of each image chunk.
138 Returns
139 -------
140 :
141 List of array chunks and their relative indices.
142 """
144 chunks = []
145 length = array.shape[4]
146 width = array.shape[3]
148 # Calculate chunk splits.
149 length_section = [0] + (int(length / chunk_size)) * [chunk_size]
150 length_splits = np.array(length_section, dtype=np.int32).cumsum()
151 width_section = [0] + (int(width / chunk_size)) * [chunk_size]
152 width_splits = np.array(width_section, dtype=np.int32).cumsum()
154 # Iterate through each chunk split.
155 for i in range(len(length_splits) - 1):
156 length_start = length_splits[i]
157 length_end = length_splits[i + 1]
159 for j in range(len(width_splits) - 1):
160 width_start = width_splits[j]
161 width_end = width_splits[j + 1]
163 # Extract chunk from full contents.
164 chunk = np.copy(array[:, :, :, length_start:length_end, width_start:width_end])
166 if np.sum(chunk) != 0:
167 chunks.append((i, j, chunk))
169 return chunks
172def flatten_array_chunk(array: np.ndarray, image_type: ImageType) -> np.ndarray:
173 """
174 Flatten array chunk along z axis.
176 When flattening to an RGBA array, each object is encoded as a unique color
177 such that the object ID = R + G*256 + B*256*256 - 1 and background pixels
178 are black (R = 0, G = 0, B = 0).
180 Parameters
181 ----------
182 array
183 Image array (dimensions in TCZYX order).
184 image_type
185 Image conversion type.
187 Returns
188 -------
189 :
190 Flattened image array.
191 """
193 array_flat = array[0, 0, :, :, :].max(axis=0)
195 if image_type == ImageType.FLAT_RGBA_BY_FRAME:
196 array_rgba = np.zeros((*array_flat.shape, 4), dtype=np.uint8)
197 array_rgba[:, :, 0] = (array_flat & 0x000000FF) >> 0
198 array_rgba[:, :, 1] = (array_flat & 0x0000FF00) >> 8
199 array_rgba[:, :, 2] = (array_flat & 0x00FF0000) >> 16
200 array_rgba[:, :, 3] = 255 # (array_flat & 0x00FF0000) >> 24
201 return array_rgba
203 if image_type in (ImageType.FLAT_BY_FRAME, ImageType.FLAT_BINARY_BY_FRAME):
204 return array_flat
206 return array