Coverage for src/precon3d/post_stitch_utils.py: 0%

1import numpy.typing as npt 

2from typing import NamedTuple 

3from pathlib import Path 

4import skimage 

5import numpy as np 

6import time 

7import typer 

8import imageio 

9import scipy 

10from PIL import Image 

11 

12Image.MAX_IMAGE_PIXELS = None 

13import tifffile 

14import re 

15from itertools import islice 

16from itertools import pairwise 

17 

18from joblib import Parallel, delayed 

19from concurrent.futures import ThreadPoolExecutor 

20 

21import precon3d.utility as ut 

22 

23# pylint: disable=wildcard-import 

24from precon3d._my_typer_cli import CustomCLIGroup, CustomCLICommand 

25 

26# CLI 

27app = typer.Typer( 

28 cls=CustomCLIGroup, 

29 no_args_is_help=True, 

30 short_help="Prepare polarized images", 

31) 

32 

33 

34def rotate_image_deg(image_path: Path, output_dir: Path): 

35 """Use XXXdeg in filename to rotate images 

36 

37 Args: 

38 image_path (Path): The path to the image file. 

39 

40 Returns: 

41 Image: The rotated image. 

42 """ 

43 

44 # Convert Path to string if not already, to use regex 

45 fname_str = str(image_path.stem) 

46 

47 # Extract the degree information from the filename using regex 

48 match = re.search(r"(\d{1,3})deg", fname_str) 

49 if match: 

50 degrees = int(match.group(1)) 

51 else: 

52 raise ValueError("No degree information found in filename.") 

53 

54 # Load the image 

55 with Image.open(image_path) as img: 

56 print(f"rotating {image_path.stem} ccw {degrees} degrees") 

57 

58 # Rotate the image. The expand flag is used to resize the output to fit the new orientation 

59 rotated_img = img.rotate(degrees, expand=False) 

60 

61 # save the rotated image or return it 

62 save_path = output_dir.joinpath(f"rotated_{image_path.name}") 

63 

64 rotated_img.save(save_path) 

65 

66 

67def rotate_image_deg_imageio(image_path: Path, degrees): 

68 # Read the image 

69 img = imageio.v3.imread(image_path) 

70 

71 rot_img = scipy.ndimage.rotate(img, degrees, reshape=True, mode="constant") 

72 

73 # # Calculate the rotation in radians 

74 # radians = np.deg2rad(degrees) 

75 

76 # # Perform the rotation 

77 # rotated_img = imageio.core.functions.rotate(img, radians, resize=True) 

78 fpath_npy = image_path.parent.joinpath(f"rotated_{image_path.stem}.npy") 

79 # fpath_tif = image_path.parent.joinpath(f"rotated_{image_path.name}") 

80 

81 # Save the rotated image 

82 # imageio.v3.imwrite(fpath_tif, rot_img) 

83 np.save(fpath_npy, rot_img) 

84 

85 

86def rotate_and_save_tiff( 

87 image_path, output_path, angle, clip_to_uint8=False, reshape=True 

88): 

89 """ 

90 Rotate a high bit-depth TIFF image and save the result. 

91 

92 Args: 

93 - image_path (str): Path to the input TIFF image. 

94 - output_path (str): Path where the rotated image will be saved. 

95 - angle (float): Rotation angle in degrees. Positive for counter-clockwise rotation. 

96 - reshape (bool): Whether to expand the output image's shape to fit the entire rotated image. 

97 """ 

98 # Read the image 

99 with tifffile.TiffFile(image_path) as tif: 

100 img_array = tif.asarray() 

101 

102 # Rotate the image 

103 rotated_img_array = scipy.ndimage.rotate( 

104 img_array, angle, reshape=reshape, order=1, mode="constant" 

105 ) 

106 

107 # Find the bounding box of non-zero values 

108 non_zero_coords = np.argwhere( 

109 rotated_img_array > 0 

110 ) # Assuming non-zero values are the content 

111 y_min, x_min = non_zero_coords.min(axis=0) 

112 y_max, x_max = ( 

113 non_zero_coords.max(axis=0) + 1 

114 ) # +1 because slice end index is exclusive 

115 

116 # Crop the image 

117 cropped_img_array = rotated_img_array[y_min:y_max, x_min:x_max] 

118 

119 # Clip values to the 0-255 range and convert to uint8 

120 if clip_to_uint8: 

121 cropped_img_array = np.clip(cropped_img_array, 0, 255).astype(np.uint8) 

122 

123 # Save the rotated image 

124 tifffile.imwrite(output_path, cropped_img_array, photometric="minisblack") 

125 

126 

127@app.command(name="rotate_images_by_filename") 

128def rotate_images_by_filename(input_dir: Path, output_dir: Path): 

129 """ 

130 rotate images in a directory using the XXXdeg in the filename 

131 """ 

132 

133 mosaic_image_file_list = ut.sorted_files(input_dir, "tif") 

134 

135 # Ensure output directory exists 

136 output_dir.mkdir(parents=True, exist_ok=True) 

137 

138 for each_image in mosaic_image_file_list: 

139 

140 # rotate_image_deg(each_image) 

141 # Convert Path to string if not already, to use regex 

142 fname_str = str(each_image.stem) 

143 # print(f"Rotating {fname_str}") 

144 

145 # Extract the degree information from the filename using regex 

146 match = re.search(r"(\d{1,3})deg", fname_str) 

147 if match: 

148 degrees = int(match.group(1)) 

149 else: 

150 raise ValueError("No degree information found in filename.") 

151 

152 print(f"Applying {degrees} deg rotation to {fname_str}") 

153 

154 # output_path = each_image.parent.joinpath("Rotated",f"rotated_{each_image.name}") 

155 output_filepath = output_dir.joinpath(f"rotated_{each_image.name}") 

156 

157 # rotate_image_deg_imageio(each_image, degrees) 

158 rotate_and_save_tiff( 

159 each_image, output_path=output_filepath, angle=degrees 

160 ) 

161 

162 

163def rotate_and_save_single_image_by_filename( 

164 input_image: Path, output_dir: Path 

165): 

166 """ 

167 rotate images in a directory using the XXXdeg in the filename 

168 """ 

169 

170 # rotate_image_deg(each_image) 

171 # Convert Path to string if not already, to use regex 

172 fname_str = str(input_image.stem) 

173 # print(f"Rotating {input_image.name}") 

174 

175 # Extract the degree information from the filename using regex 

176 match = re.search(r"(\d{1,3})deg", fname_str) 

177 if match: 

178 degrees = int(match.group(1)) 

179 else: 

180 raise ValueError("No degree information found in filename.") 

181 

182 print(f"Applying {degrees} deg rotation to {input_image.name}") 

183 

184 # output_path = each_image.parent.joinpath("Rotated",f"rotated_{each_image.name}") 

185 output_filepath = output_dir.joinpath(f"rotated_{input_image.name}") 

186 

187 # rotate_image_deg_imageio(each_image, degrees) 

188 rotate_and_save_tiff( 

189 input_image, output_path=output_filepath, angle=degrees 

190 ) 

191 

192 

193def rotate_save_pad_single_image_by_filename( 

194 input_image: Path, max_image_extent: np.ndarray, output_dir: Path 

195): 

196 """ 

197 rotate images in a directory using the XXXdeg in the filename 

198 """ 

199 

200 # rotate_image_deg(each_image) 

201 # Convert Path to string if not already, to use regex 

202 fname_str = str(input_image.stem) 

203 # print(f"Rotating {input_image.name}") 

204 

205 # Extract the degree information from the filename using regex 

206 match = re.search(r"(\d{1,3})deg", fname_str) 

207 if match: 

208 degrees = int(match.group(1)) 

209 else: 

210 raise ValueError("No degree information found in filename.") 

211 

212 # output_path = each_image.parent.joinpath("Rotated",f"rotated_{each_image.name}") 

213 output_filepath = output_dir.joinpath( 

214 f"{input_image.stem}_pad_rot{input_image.suffix}" 

215 ) 

216 

217 # rotate_image_deg_imageio(each_image, degrees) 

218 # rotate_and_save_tiff(input_image, output_path=output_filepath,angle=degrees) 

219 # Read the image 

220 with tifffile.TiffFile(input_image) as tif: 

221 img_array = tif.asarray() 

222 

223 # Rotate the image 

224 rotated_img_array = scipy.ndimage.rotate( 

225 img_array, degrees, reshape=True, order=1, mode="constant" 

226 ) 

227 

228 # Ensure the image does not exceed the maximum dimensions 

229 cropped_img_array = rotated_img_array[ 

230 : min(rotated_img_array.shape[0], max_image_extent[0]), 

231 : min(rotated_img_array.shape[1], max_image_extent[1]), 

232 ] 

233 

234 # # Find the bounding box of non-zero values 

235 # non_zero_coords = np.argwhere(rotated_img_array > 0) # Assuming non-zero values are the content 

236 # y_min, x_min = non_zero_coords.min(axis=0) 

237 # y_max, x_max = non_zero_coords.max(axis=0) + 1 # +1 because slice end index is exclusive 

238 

239 # # Crop the image 

240 # cropped_img_array = rotated_img_array[y_min:y_max, x_min:x_max] 

241 

242 # pad_height = max_image_extent[0] - cropped_img_array.shape[0] 

243 # pad_width = max_image_extent[1] - cropped_img_array.shape[1] 

244 

245 # # Calculate padding for height and width 

246 # pad_top = pad_height // 2 

247 # pad_bottom = pad_height - pad_top 

248 # pad_left = pad_width // 2 

249 # pad_right = pad_width - pad_left 

250 

251 # Calculate necessary padding 

252 pad_height = max(0, max_image_extent[0] - cropped_img_array.shape[0]) 

253 pad_width = max(0, max_image_extent[1] - cropped_img_array.shape[1]) 

254 

255 # Calculate padding for height and width 

256 pad_top = pad_height // 2 

257 pad_bottom = pad_height - pad_top 

258 pad_left = pad_width // 2 

259 pad_right = pad_width - pad_left 

260 

261 # Ensure padding does not exceed the image dimensions 

262 pad_top = min(pad_top, max_image_extent[0]) 

263 pad_bottom = min(pad_bottom, max_image_extent[0] - pad_top) 

264 pad_left = min(pad_left, max_image_extent[1]) 

265 pad_right = min(pad_right, max_image_extent[1] - pad_left) 

266 

267 # Apply padding 

268 padded_image = np.pad( 

269 cropped_img_array, 

270 ((pad_top, pad_bottom), (pad_left, pad_right)), 

271 mode="constant", 

272 constant_values=0, 

273 ) 

274 

275 # Save the padded image 

276 skimage.io.imsave(output_filepath, padded_image, check_contrast=False) 

277 

278 print( 

279 f"Saved {output_filepath.name} after a {degrees} deg rotation and padded to {max_image_extent}" 

280 ) 

281 

282 

283def pad_images_to_max_extent(input_dir: Path, output_dir: Path): 

284 """Pad all images to the same size, the size is the max extent found in the images 

285 

286 Args: 

287 input_dir (Path): Directory containing input images. 

288 output_dir (Path): Directory where padded images will be saved. 

289 """ 

290 filepaths_list = ut.sorted_files( 

291 directory=input_dir, file_extension=".tif" 

292 ) 

293 n_images = len(filepaths_list) 

294 print(f"Found {n_images} images") 

295 

296 # Collect all images and their shapes 

297 all_images_shape = [] 

298 for item in filepaths_list: 

299 each_image = skimage.io.imread(item) 

300 all_images_shape.append(each_image.shape) 

301 

302 # Determine the maximum extent in each dimension 

303 max_image_extent = np.max(all_images_shape, axis=0) 

304 print(f"\tPadding images to {max_image_extent} pixels") 

305 

306 # Ensure output directory exists 

307 output_dir.mkdir(parents=True, exist_ok=True) 

308 

309 # Pad and save images 

310 for each_image_path in filepaths_list: 

311 each_image = skimage.io.imread(each_image_path) 

312 pad_height = max_image_extent[0] - each_image.shape[0] 

313 pad_width = max_image_extent[1] - each_image.shape[1] 

314 

315 # Calculate padding for height and width 

316 pad_top = pad_height // 2 

317 pad_bottom = pad_height - pad_top 

318 pad_left = pad_width // 2 

319 pad_right = pad_width - pad_left 

320 

321 # Apply padding 

322 padded_image = np.pad( 

323 each_image, 

324 ((pad_top, pad_bottom), (pad_left, pad_right)), 

325 mode="constant", 

326 constant_values=0, 

327 ) 

328 

329 # Save the padded image 

330 tif_out_path = output_dir.joinpath(each_image_path.name) 

331 skimage.io.imsave(tif_out_path, padded_image, check_contrast=False) 

332 

333 # print("Padding and saving completed.") 

334 

335 

336def image_max_extent(filepaths_list: list[Path]): 

337 """ 

338 Retrieve the maximum image dimensions from a list of file paths. 

339 

340 Leverages PIL Image package for quick access to image width and height. 

341 """ 

342 n_images = len(filepaths_list) 

343 print(f"Found {n_images} images") 

344 

345 def get_image_size(filepath): 

346 with Image.open(filepath) as img: 

347 return img.size # Returns (width, height) 

348 

349 # Use ThreadPoolExecutor to parallelize the retrieval of image sizes 

350 with ThreadPoolExecutor() as executor: 

351 all_image_sizes = list(executor.map(get_image_size, filepaths_list)) 

352 

353 # Convert (width, height) to (height, width) and find max dimensions 

354 all_image_sizes = [(height, width) for width, height in all_image_sizes] 

355 max_image_extent = np.max(all_image_sizes, axis=0) 

356 print(f"Padding images to {max_image_extent} pixels") 

357 

358 return max_image_extent 

359 

360 

361def pad_and_save_image_equal( 

362 filename: Path, max_image_extent: np.ndarray, output_dir: Path 

363): 

364 """To come""" 

365 

366 input_image = skimage.io.imread(filename) 

367 pad_height = max_image_extent[0] - input_image.shape[0] 

368 pad_width = max_image_extent[1] - input_image.shape[1] 

369 

370 # Calculate padding for height and width 

371 pad_top = pad_height // 2 

372 pad_bottom = pad_height - pad_top 

373 pad_left = pad_width // 2 

374 pad_right = pad_width - pad_left 

375 

376 # Apply padding 

377 if input_image.ndim == 3 and input_image.shape[2] == 3: 

378 padded_image = np.pad( 

379 input_image, 

380 ( 

381 (pad_top, pad_bottom), 

382 (pad_left, pad_right), 

383 (0, 0), 

384 ), # No padding for the channel dimension 

385 mode="constant", 

386 constant_values=0, 

387 ) # grayscale 

388 else: 

389 padded_image = np.pad( 

390 input_image, 

391 ((pad_top, pad_bottom), (pad_left, pad_right)), 

392 mode="constant", 

393 constant_values=0, 

394 ) 

395 

396 # Save the padded image 

397 output_filename = f"padded_{filename.name}" 

398 tif_out_path = output_dir.joinpath(output_filename) 

399 skimage.io.imsave(tif_out_path, padded_image, check_contrast=False) 

400 

401 

402def filter_and_sort_file_paths( 

403 file_paths: list[Path], keyword: str 

404) -> list[Path]: 

405 """ 

406 Filters a list of file paths to include only those that contain a specific directory keyword and sorts the result. 

407 

408 Args: 

409 file_paths (list of str): The list of file paths to filter. 

410 keyword (str): The directory keyword to filter by. 

411 

412 Returns: 

413 list of Path: A sorted list containing only the file paths that include the keyword. 

414 """ 

415 # Convert strings to Path objects and filter 

416 filtered_paths = [ 

417 Path(path) for path in file_paths if keyword in Path(path).parts 

418 ] 

419 

420 # Sort the filtered paths 

421 sorted_filtered_paths = sorted(filtered_paths, key=lambda x: x.as_posix()) 

422 

423 return sorted_filtered_paths 

424 

425 

426def append_deg_suffix_and_rename(filenames: list[Path], directory_path: Path): 

427 # Define the degree suffixes in a list 

428 deg_suffixes = [ 

429 "000deg", 

430 "015deg", 

431 "030deg", 

432 "045deg", 

433 "060deg", 

434 "075deg", 

435 "090deg", 

436 "105deg", 

437 "120deg", 

438 "135deg", 

439 "150deg", 

440 "165deg", 

441 "180deg", 

442 ] 

443 

444 # Process each filename 

445 updated_filenames = [] 

446 for i, filename in enumerate(filenames): 

447 # Calculate the modulo 13 of the index (i + 1) to match the desired mapping 

448 mod_value = i % 13 

449 suffix = deg_suffixes[mod_value] 

450 

451 # Remove existing duplicate suffixes 

452 stem = filename.stem 

453 while stem.endswith(suffix): 

454 stem = stem[ 

455 : -(len(suffix) + 1) 

456 ] # Remove the suffix and underscore 

457 

458 # Append the corresponding degree suffix to the cleaned filename 

459 new_filename = f"{stem}_{suffix}{filename.suffix}" 

460 updated_filenames.append(new_filename) 

461 

462 # Full path for old and new filenames 

463 old_file_path = directory_path / filename 

464 new_file_path = directory_path / new_filename 

465 

466 # Rename the file 

467 old_file_path.rename(new_file_path) 

468 

469 return updated_filenames 

470 

471 

472@app.command( 

473 cls=CustomCLICommand, 

474 name="rotate_images_by_filename_parallel", 

475 short_help="rotate images by filename in parallel", 

476) 

477def rotate_images_by_filename_parallel( 

478 input_dir: Path, output_dir: Path, ncores: int = 4 

479): 

480 """ 

481 rotate images in a directory using the XXXdeg in the filename 

482 """ 

483 

484 image_file_list = ut.sorted_files(input_dir, "tif") 

485 

486 # Ensure output directory exists 

487 output_dir.mkdir(parents=True, exist_ok=True) 

488 

489 # rotate and save images, parallel 

490 # Parallel(n_jobs=ncores)(delayed(rotate_and_save_single_image_by_filename)(each_image_path, output_dir) for each_image_path in image_file_list) 

491 # Parallel(n_jobs=ncores)(delayed(rotate_image_deg)(each_image_path, output_dir) for each_image_path in image_file_list) 

492 Parallel(n_jobs=ncores)( 

493 delayed(rotate_and_save_single_image_by_filename)( 

494 each_image_path, output_dir 

495 ) 

496 for each_image_path in image_file_list 

497 ) 

498 

499 print("Image rotations completed.") 

500 

501 

502@app.command( 

503 cls=CustomCLICommand, 

504 name="pad_images_to_max_extent_parallel", 

505 short_help="pad images to maximum extent in parallel", 

506) 

507def pad_images_to_max_extent_parallel( 

508 input_dir: Path, output_dir: Path, ncores: int = 4 

509): 

510 """Pad all images to the same size, the size is the max extent determined from all the images 

511 

512 Args: 

513 input_dir (Path): Directory containing input images. 

514 output_dir (Path): Directory where padded images will be saved. 

515 """ 

516 # filepaths_list = ut.sorted_files( 

517 # directory=input_dir, file_extension=".tif" 

518 # ) 

519 

520 filepaths_list = ut.sorted_files( 

521 directory=input_dir, file_extension=".tif" 

522 ) 

523 

524 max_image_extent = image_max_extent(filepaths_list) 

525 

526 # Ensure output directory exists 

527 output_dir.mkdir(parents=True, exist_ok=True) 

528 

529 # Pad and save images, parallel 

530 

531 Parallel(n_jobs=ncores)( 

532 delayed(pad_and_save_image_equal)( 

533 each_image_path, max_image_extent, output_dir 

534 ) 

535 for each_image_path in filepaths_list 

536 ) 

537 

538 print("Padding and saving completed.") 

539 

540 

541def combined_image_padding_and_rotation_parallel( 

542 input_dir: Path, 

543 max_image_extent: np.ndarray, 

544 output_dir: Path, 

545 ncores: int = 4, 

546): 

547 """To come""" 

548 

549 filepaths_list = ut.sorted_files( 

550 directory=input_dir, file_extension=".tif" 

551 ) 

552 

553 # Ensure output directory exists 

554 output_dir.mkdir(parents=True, exist_ok=True) 

555 

556 # Pad and save images, parallel 

557 Parallel(n_jobs=ncores)( 

558 delayed(rotate_save_pad_single_image_by_filename)( 

559 each_image_path, max_image_extent, output_dir 

560 ) 

561 for each_image_path in filepaths_list 

562 ) 

563 

564 print("Padding, rotation, and saving completed.") 

565 

566 

567@app.callback() 

568def callback(): 

569 """ 

570 precon3d.pol_prep can prepare polarized images by padding and rotating in a parallel operation. Colorization also available. 

571 """ 

572 

573 

574import matplotlib.pyplot as plt 

575 

576 

577def visualize_rgb(rgb_image: npt.NDArray): 

578 """plot the images side by side""" 

579 

580 H, W, _ = rgb_image.shape 

581 

582 # 1) Normalize each component map to [0,1] for display 

583 C_norm = np.empty_like(rgb_image, dtype=float) 

584 for i in range(3): 

585 comp = rgb_image[:, :, i] 

586 lo, hi = comp.min(), comp.max() 

587 if hi > lo: 

588 C_norm[:, :, i] = (comp - lo) / (hi - lo) 

589 else: 

590 C_norm[:, :, i] = comp 

591 

592 # 2) Build the RGB composite 

593 rgb_composite = C_norm # channel 0→R, 1→G, 2→B 

594 

595 # 3) Build three “pure‐R/G/B” versions of each component 

596 colored = [] 

597 for i in range(3): 

598 # zero‐plane 

599 zero = np.zeros((H, W), dtype=float) 

600 if i == 0: 

601 # comp 0 in red channel 

602 rgb = np.stack([C_norm[:, :, 0], zero, zero], axis=-1) 

603 elif i == 1: 

604 # comp 1 in green 

605 rgb = np.stack([zero, C_norm[:, :, 1], zero], axis=-1) 

606 else: 

607 # comp 2 in blue 

608 rgb = np.stack([zero, zero, C_norm[:, :, 2]], axis=-1) 

609 colored.append(rgb) 

610 

611 # 4) Plot the 1×4 montage 

612 fig, axes = plt.subplots(1, 4, figsize=(4 * 4, 4)) 

613 

614 # Panel 1: RGB composite 

615 ax = axes[0] 

616 ax.imshow(rgb_composite, vmin=0, vmax=1) 

617 ax.set_title("Composite (RGB)") 

618 ax.axis("off") 

619 

620 # Panels 2–4: individual comps in R/G/B 

621 titles = ["Component 1 → Red", "Component 2 → Green", "Component 3 → Blue"] 

622 for ax, img, title in zip(axes[1:], colored, titles): 

623 ax.imshow(img, vmin=0, vmax=1) 

624 ax.set_title(title) 

625 ax.axis("off") 

626 

627 plt.tight_layout() 

628 plt.show() 

629 

630 

631def visualize_spectra(spectra: npt.NDArray): 

632 """plot the spectra, color as RGB""" 

633 

634 n_channels = spectra.shape[0] 

635 channels = np.arange(n_channels) # or your actual channel‐wavelength array 

636 

637 plt.figure(figsize=(6, 4)) 

638 plt.plot(channels, spectra[:, 0], color="red", lw=2, label="Comp 1") 

639 plt.plot(channels, spectra[:, 1], color="green", lw=2, label="Comp 2") 

640 plt.plot(channels, spectra[:, 2], color="blue", lw=2, label="Comp 3") 

641 

642 plt.xlabel("Channel") 

643 plt.ylabel("Intensity") 

644 plt.title("Recovered Spectra") 

645 plt.legend() 

646 plt.grid(True) 

647 plt.tight_layout() 

648 plt.show() 

649 

650 

651import os 

652from skimage import io 

653from skimage.util import img_as_float 

654 

655 

656def load_pol_images( 

657 img_dir: Path, file_extension: str = ".tif" 

658) -> npt.NDArray: 

659 """load and reshape images stack data from directory""" 

660 

661 # Create an ImageCollection 

662 # image_collection = io.ImageCollection(os.path.join(img_dir, '*.tif')) # Adjust the file extension as needed 

663 # image_collection = io.ImageCollection(ut.sorted_files(img_dir, file_extension)) 

664 

665 # Optionally, convert images to float format 

666 # image_collection_float = [img_as_float(image) for image in image_collection] 

667 

668 # Example: Print the number of images loaded and their shapes 

669 # print(f'Number of images loaded: {len(image_collection)}') 

670 # for i, img in enumerate(image_collection_float): 

671 # print(f'Image {i}: shape {img.shape}') 

672 

673 image_stack_paths = ut.sorted_files(img_dir, file_extension) 

674 

675 # Load the shading images into a list, using imageio for image reading 

676 imagestack = [ 

677 imageio.v3.imread(each_image) for each_image in image_stack_paths 

678 ] 

679 

680 # Reshape the loaded images into a 2D array where each row represents a flattened image 

681 reshaped_images = np.array([img_as_float(img) for img in imagestack]) 

682 

683 img = np.asarray(reshaped_images) 

684 

685 img = np.transpose(img, (1, 2, 0)) 

686 

687 return img 

688 

689 

690from sklearn.decomposition import PCA, NMF 

691 

692 

693def varimax(Phi, gamma=1.0, q=20, tol=1e-6): 

694 """ 

695 Perform varimax (orthogonal) rotation on the loadings matrix Phi. 

696 Input: 

697 Phi = (p, k) loadings 

698 gamma = 1.0 for classic varimax 

699 q = max iterations 

700 tol = convergence tolerance 

701 Returns: 

702 rotated loadings, rotation matrix R 

703 """ 

704 p, k = Phi.shape 

705 R = np.eye(k) 

706 d = 0 

707 for i in range(q): 

708 Lambda = Phi.dot(R) 

709 # the varimax update 

710 u, s, vh = np.linalg.svd( 

711 Phi.T.dot( 

712 np.power(Lambda, 3) 

713 - (gamma / p) * Lambda.dot(np.diag(np.sum(Lambda**2, axis=0))) 

714 ) 

715 ) 

716 R = u.dot(vh) 

717 d_new = np.sum(s) 

718 if d != 0 and (d_new - d) < tol: 

719 break 

720 d = d_new 

721 return Phi.dot(R), R 

722 

723 

724@app.command( 

725 cls=CustomCLICommand, 

726 name="colorize_pca_varimax", 

727 short_help="Colorize with PCA", 

728) 

729def colorize_pca_varimax(img_dir: Path): 

730 """colorize the polarized iamge stack with pca and varimax""" 

731 

732 # cast to Path object if not already a Path object 

733 if not isinstance(img_dir, Path): 

734 img_dir = Path(img_dir) 

735 

736 img = load_pol_images(img_dir) 

737 

738 H, W, C = img.shape 

739 # reshape into (n_pixels, n_channels) 

740 X = img.reshape(-1, 13) # shape = (H*W, 13) 

741 

742 n_components = 3 # for RGB 

743 

744 # 3a. fit PCA 

745 pca = PCA(n_components=n_components) 

746 scores = pca.fit_transform(X) # (H*W, 3) principal component scores 

747 loadings = pca.components_.T # (13, 3) unrotated loadings 

748 

749 # 3b. rotate loadings 

750 rot_loadings, R = varimax(loadings) # (13, 3), rotation R 

751 # rotated scores 

752 rot_scores = X.dot(rot_loadings) # (H*W, 3) 

753 

754 # reshape back to images 

755 score_imgs = rot_scores.reshape(H, W, n_components) 

756 

757 visualize_rgb(score_imgs) 

758 visualize_spectra(rot_loadings) 

759 

760 # Normalize before saving 

761 score_imgs_norm = np.empty_like(score_imgs, dtype=float) 

762 for i in range(3): 

763 comp = score_imgs[:, :, i] 

764 lo, hi = comp.min(), comp.max() 

765 if hi > lo: 

766 score_imgs_norm[:, :, i] = (comp - lo) / (hi - lo) 

767 else: 

768 score_imgs_norm[:, :, i] = comp 

769 

770 # io.imsave(img_dir.parent.joinpath('rgb_pca_image.png'), (score_imgs * 255).astype(np.uint8), check_contrast=False) 

771 imageio.v3.imwrite( 

772 img_dir.parent.joinpath(f"{img_dir.name}_rgb_pca_image.tif"), 

773 (score_imgs_norm * 255).astype(np.uint8), 

774 ) 

775 

776 

777def mcr_als(X, C_init, S_init, max_iter=100, tol=1e-6): 

778 """ 

779 Basic non‐negative MCR‐ALS. 

780 X : (n_samples, n_vars) 

781 C_init : (n_samples, n_comp) 

782 S_init : (n_vars, n_comp) 

783 """ 

784 C = C_init.copy() 

785 S = S_init.copy() 

786 for it in range(max_iter): 

787 S_old = S.copy() 

788 

789 # 1) update spectra: X = C @ S.T ⇒ S.T = pinv(C) @ X 

790 S = (np.linalg.pinv(C) @ X).T 

791 S[S < 0] = 0 

792 

793 # 2) update concentrations: X = C @ S.T ⇒ C = X @ pinv(S.T) 

794 C = X @ np.linalg.pinv(S.T) 

795 C[C < 0] = 0 

796 

797 # 3) convergence on S 

798 delta = np.linalg.norm(S - S_old) / np.linalg.norm(S_old) 

799 if delta < tol: 

800 print(f"MCR‐ALS converged in {it+1} iter (Δ={delta:.2e})") 

801 break 

802 

803 return C, S 

804 

805 

806def mcr_als_lstsq(X, C_init, S_init, max_iter=100, tol=1e-6): 

807 C = C_init.copy() 

808 S = S_init.copy() 

809 for it in range(max_iter): 

810 S_old = S.copy() 

811 

812 # 1) solve C @ S.T = X ⇒ S.T = lstsq(C, X) 

813 S = np.linalg.lstsq(C, X, rcond=None)[0].T 

814 S[S < 0] = 0 

815 

816 # 2) solve S @ C.T = X.T ⇒ C.T = lstsq(S, X.T) 

817 C = np.linalg.lstsq(S, X.T, rcond=None)[0].T 

818 C[C < 0] = 0 

819 

820 delta = np.linalg.norm(S - S_old) / np.linalg.norm(S_old) 

821 if delta < tol: 

822 print(f"MCR‐ALS converged in {it+1} iter (Δ={delta:.2e})") 

823 break 

824 

825 return C, S 

826 

827 

828@app.command( 

829 cls=CustomCLICommand, 

830 name="colorize_mcr_als", 

831 short_help="Colorize with MCR", 

832) 

833def colorize_mcr_als(img_dir: Path): 

834 """colorize the polarized iamge stack with mcr and als""" 

835 

836 # cast to Path object if not already a Path object 

837 if not isinstance(img_dir, Path): 

838 img_dir = Path(img_dir) 

839 

840 img = load_pol_images(img_dir) 

841 

842 H, W, C = img.shape 

843 # reshape into (n_pixels, n_channels) 

844 X = img.reshape(-1, 13) # shape = (H*W, 13) 

845 

846 n_components = 3 # for RGB 

847 

848 # 4a. initialize with NMF (non‐negative matrix factorization) 

849 nmf = NMF( 

850 n_components=n_components, 

851 init="random", 

852 random_state=0, 

853 max_iter=500, 

854 tol=1e-4, 

855 ) 

856 C_init = nmf.fit_transform(X) # (H*W, 3) initial concentration profiles 

857 S_init = nmf.components_ # (3, 13) initial spectra 

858 

859 # transpose S_init so spectra are columns (13, 3) 

860 S_init = S_init.T 

861 

862 C_mcr, S_mcr = mcr_als(X, C_init, S_init) 

863 

864 # reshape concentrations back to image form 

865 C_imgs = C_mcr.reshape(H, W, n_components) 

866 

867 visualize_rgb(C_imgs) 

868 visualize_spectra(S_mcr) 

869 

870 # io.imsave(img_dir.parent.joinpath('rgb_mcr_image.png'), (C_imgs * 255).astype(np.uint8), check_contrast=False) 

871 imageio.v3.imwrite( 

872 img_dir.parent.joinpath(f"{img_dir.name}_rgb_mcr_image.tif"), 

873 (C_imgs * 255).astype(np.uint8), 

874 ) 

875 

876 

877def apply_loading(img_dir: Path, loadings: npt.NDArray): 

878 """use a known loading to color polarized images""" 

879 

880 # cast to Path object if not already a Path object 

881 if not isinstance(img_dir, Path): 

882 img_dir = Path(img_dir) 

883 

884 img = load_pol_images(img_dir) 

885 

886 H, W, C = img.shape 

887 # reshape into (n_pixels, n_channels) 

888 X = img.reshape(-1, 13) # shape = (H*W, 13) 

889 

890 rng = np.random.default_rng()