VisionLanguageGroup/MicroscopyMatching

1import os2from glob import glob3from pathlib import Path4from natsort import natsorted5from PIL import Image6import numpy as np7import tifffile8import skimage.io as io9import torchvision.transforms as T10import cv211from tqdm import tqdm12from models.tra_post_model.utils import normalize_01, normalize13IMG_SIZE = 51214 15def _load_tiffs(folder: Path, dtype=None):16    """Load a sequence of tiff files from a folder into a 3D numpy array."""17    images = glob(str(folder / "*.tif"))18    test_data = tifffile.imread(images[0])19    if len(test_data.shape) == 3:20        turn_gray = True21    else:22        turn_gray = False23    end_frame = len(images)24    if not turn_gray:25        x = np.stack([26            tifffile.imread(f).astype(dtype)27            for f in tqdm(28                sorted(folder.glob("*.tif"))[0 : end_frame : 1],29                leave=False,30                desc=f"Loading [0:{end_frame}]",31            )32        ])33    else:34        x = []35        for f in tqdm(36            sorted(folder.glob("*.tif"))[0 : end_frame : 1],37            leave=False,38            desc=f"Loading [0:{end_frame}]",39        ):40            img = tifffile.imread(f).astype(dtype)41            if img.ndim == 3:42                if img.shape[-1] > 3:43                    img = img[..., :3]44                img = (0.299 * img[..., 0] + 0.587 * img[..., 1] + 0.114 * img[..., 2])45            x.append(img)46        x = np.stack(x)47    return x48 49 50def load_track_images(file_dir):51    52    def find_tif_dir(root_dir):53        tif_files = []54        for dirpath, _, filenames in os.walk(root_dir):55            if '__MACOSX' in dirpath:56                continue57            for f in filenames:58                if f.lower().endswith('.tif'):59                    tif_files.append(os.path.join(dirpath, f))60        return tif_files61 62    tif_dir = find_tif_dir(file_dir)63    print(f"Found {len(tif_dir)} tif images in {file_dir}")64    print(f"First 5 tif images: {tif_dir[:5]}")65    assert len(tif_dir) > 0, f"No tif images found in {file_dir}"66    images = natsorted(tif_dir)67    imgs = []68    imgs_raw = []69    images_stable = []70    # load images for seg and track71    for img_path in tqdm(images, desc="Loading images"):72        img = tifffile.imread(img_path)73        img_raw = io.imread(img_path)74    75        if img.dtype == 'uint16':76            img = ((img - img.min()) / (img.max() - img.min() + 1e-6) * 255).astype(np.uint8)77            img = np.stack([img] * 3, axis=-1)78            w, h = img.shape[1], img.shape[0]79        else:80            img = Image.open(img_path).convert("RGB")81            w, h = img.size82 83        img = T.Compose([84            T.ToTensor(),85            T.Resize((IMG_SIZE, IMG_SIZE)),86        ])(img)87 88        image_stable = img - 0.589        img = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)90 91 92        imgs.append(img)93        imgs_raw.append(img_raw)94        images_stable.append(image_stable)95 96    height = h97    width = w98    imgs = np.stack(imgs, axis=0)99    imgs_raw = np.stack(imgs_raw, axis=0)100    images_stable = np.stack(images_stable, axis=0)101 102    # track data103    imgs_ = _load_tiffs(Path(file_dir), dtype=np.float32)104    imgs_01 = np.stack([105                normalize_01(_x) for _x in tqdm(imgs_, desc="Normalizing", leave=False)106            ])107    imgs_ = np.stack([108                normalize(_x) for _x in tqdm(imgs_, desc="Normalizing", leave=False)109            ])110 111    return imgs, imgs_raw, images_stable, imgs_, imgs_01, height, width112 113