Visualize a TEM section with each subtile marked as resin or not resin.

resin_mask.py

"""
Extracts tissue from an overview montage and calculates the centroid.

python resin_mask.py .../bladeseq-2025.02.03-11.41.01/s021-2025.02.03-11.41.01/

To install dependencies:
pip install simplejpeg tqdm pyspng-seunglab numpy scipy connected-components-3d fill_voids tinybrain edt opencv-python -U 
"""
from typing import List, Dict, Union, Optional, Tuple

import csv
import io
import re
import sys
import os

import numpy as np
import numpy.typing as npt
import cc3d
import simplejpeg
import pyspng
import scipy.ndimage
import scipy.signal
import tinybrain
import microviewer
import edt
import time
import fill_voids
import fastremap
from tqdm import tqdm
import scipy.signal
import cv2 as cv

from PIL import Image

import matplotlib.pyplot as plt

StageInfo_t = List[Dict[str,Union[int,float]]]

def is_resin(img:npt.NDArray[np.uint8]) -> bool:
	ds = tinybrain.downsample_with_averaging(img, (2,2,1), num_mips=3)
	img = ds[-1]
	hist, bin_edges = np.histogram(img.ravel(), bins=20)
	peak_idxs, bounds = scipy.signal.find_peaks(hist, height=500, threshold=4000)

	if len(peak_idxs) != 1:
		return False

	mean = np.mean(img)
	if mean <= 185:
		return False

	stdev = np.std(img)
	if stdev >= 11:
		return False

	edges = cv.Canny(img, threshold1=120, threshold2=200)
	edges = cc3d.dust(edges, threshold=35)

	return not np.any(edges)

def decode_image(filename):
	with open(filename, "rb") as f:
		binary = f.read()

	_, ext = os.path.splitext(filename)
	if ext == ".png":
		info = pyspng.header(binary)
		img = pyspng.load(binary).reshape(-1)
		img = img.astype(np.uint8, copy=False)
		return img.reshape([
			info["width"], info["height"] 
		], order="F")
	elif ext in (".jpg", ".jpeg"):
		return simplejpeg.decode_jpeg(
			binary, 
			colorspace="gray",
		)
	elif ext == ".bmp":
		img = Image.open(io.BytesIO(binary))
		return np.asarray(img, dtype=np.uint8)
	else:
		raise ValueError(f"unable to decode image of unknown type: {filename}")

def read_stage_csv(section_path) -> StageInfo_t:
    """
    Example data:
    tile_id, stage_x_nm, stage_y_nm, x_relroi_nm, y_relroi_nm
    0, 338700.639, 109345.286, 329885.172, 109961.724
    1, 283719.777, 109345.286, 274904.310, 109961.724
    """
    filepath = os.path.join(section_path, "metadata/stage_positions.csv")

    with open(filepath, "rt") as f:
    	stage_data = io.StringIO(f.read())

    stage_csv = []
    for row in csv.DictReader(stage_data):
        tile_id = int(row["tile_id"])
        row = { 
            k.strip():float(v.strip())
            for k,v in row.items() 
        }
        row["tile_id"] = tile_id
        stage_csv.append(row)
    
    return stage_csv

def compute_supertile_map(stage_info:StageInfo_t) -> np.ndarray:
    # Determine the dimensions of the 2D array
    rank_x = np.unique([ row["stage_x_nm"] for row in stage_info ])
    rank_y = np.unique([ row["stage_y_nm"] for row in stage_info ])

    rank_x = { v:i for i,v in enumerate(rank_x) }
    rank_y = { v:i for i,v in enumerate(rank_y) }    
    
    arr = np.full((len(rank_y)+1, len(rank_x)+1), None)

    for row in stage_info:
        arr[rank_y[row["stage_y_nm"]], rank_x[row["stage_x_nm"]]] = row["tile_id"]
    
    # Reverse the order of the rows in the array
    supertile_map = arr[::-1]
    
    return supertile_map

def create_resin_mask(section_path):
	stage_data = read_stage_csv(section_path)
	tile_map = compute_supertile_map(stage_data)

	tile_id2xy = {}
	for i in range(tile_map.shape[0]):
		for j in range(tile_map.shape[1]):
			tile_id2xy[tile_map[i,j]] = (j,i)

	ds_subtile_width = 375 # 6000 / 2^4

	img = np.zeros([ 3 * tile_map.shape[1] * ds_subtile_width, 3 * tile_map.shape[0] * ds_subtile_width ], dtype=np.uint8)
	resin_subtile_mask = np.zeros(3 * np.array(tile_map.shape[::-1]), dtype=bool)
	resin_mask = np.zeros(3 * ds_subtile_width * np.array(tile_map.shape[::-1]), dtype=bool)

	TILE_REGEXP = re.compile(r'tile_(\d+)_(\d).bmp')

	subtiles_path = os.path.join(section_path, "subtiles")

	subtile2xy = [
		(1,1),
		(2,1),
		(2,2),
		(1,2),
		(0,2),
		(0,1),
		(0,0),
		(1,0),
		(2,0)
	]

	paths = os.listdir(subtiles_path)
	paths.sort()

	for fname in tqdm(paths):
		m = TILE_REGEXP.match(fname)
		if m is None:
			continue
		supertile_id, subtile_id = m.groups()
		supertile_id = int(supertile_id)
		subtile_id = int(subtile_id)

		x,y = tile_id2xy[supertile_id]
		dx,dy = subtile2xy[subtile_id]
		single_subtile = decode_image(os.path.join(subtiles_path, fname))

		X = 3*x + dx
		Y = 3*y + dy

		resin_subtile_mask[X,Y] = is_resin(single_subtile)

		single_subtile = tinybrain.downsample_with_averaging(single_subtile, factor=(2,2,1), num_mips=4)[-1]
		img[
			X*ds_subtile_width:(X+1)*ds_subtile_width,
			Y*ds_subtile_width:(Y+1)*ds_subtile_width
		] = single_subtile.T

		resin_mask[
			X*ds_subtile_width:(X+1)*ds_subtile_width,
			Y*ds_subtile_width:(Y+1)*ds_subtile_width
		] = resin_subtile_mask[X,Y]

	
	img = tinybrain.downsample_with_averaging(img, factor=(2,2,1), num_mips=2)[-1]
	resin_mask = tinybrain.downsample_segmentation(resin_mask.view(np.uint8), factor=(2,2,1), num_mips=2)[-1]

	microviewer.hyperview(img, resin_mask)

if __name__ == "__main__":
	create_resin_mask(sys.argv[1])