TEM montage overview ROI centroid finder

tissue_finder.py

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

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

import csv
import io
import re
import sys
import os.path

import numpy as np
import cc3d
import simplejpeg
import pyspng
import scipy.ndimage
import tinybrain
import microviewer
import edt

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

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",
		)
	else:
		raise ValueError("unable to decode image of unknown type")

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 view_supertile_map(montage_image, supertile_map):
	supertile_map_img = np.zeros(montage_image.shape, dtype=np.uint16)
	for x in range(supertile_map.shape[0]-1):
		for y in range(supertile_map.shape[1]-1):
			supertile_map_img[x*supertile_width//2:(x+1)*supertile_width//2, y*supertile_height//2:(y+1)*supertile_width//2] = supertile_map[x+1,y]

	microviewer.hyperview(montage_image, supertile_map_img)

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 build_montage(section_path:str, stage_info:StageInfo_t) -> np.ndarray:
	supertile_map = compute_supertile_map(stage_info)

	random_tile = "dne"
	for fname in os.listdir(section_path):
		if 'supertile' in fname:
			random_tile = os.path.join(section_path, fname)
			break

	with open(random_tile, "rb") as f:
		binary = f.read()
	py, px, _, _ = simplejpeg.decode_jpeg_header(binary)

	tx, ty = supertile_map.shape

	sx = tx * px
	sy = ty * py

	montage_image = np.zeros([sx,sy], dtype=np.uint8, order="F")

	for y in range(ty):
		for x in range(tx):
			if supertile_map[x,y] is None:
				continue
			section_id = supertile_map[x,y]
			filename = os.path.join(section_path, f"supertile_{section_id}.jpg")
			try:
				supertile = decode_image(filename)
			except FileNotFoundError:
				continue
			montage_image[x*px:(x+1)*px, y*py:(y+1)*py] = supertile[:,:,0]

	montage_image = np.asfortranarray(montage_image[px:,:-py])

	return { 
		"montage_image": montage_image,
		"supertile_map": supertile_map,
		"supertile_width": px, 
		"supertile_height": py,
	}

def find_tissue_roi_kernel(
	montage_image:np.ndarray, 
	supertile_width:int, 
	supertile_height:int, 
	supertile_map:np.ndarray,
	visualize_roi:bool,
) -> int:
	"""
	Given a path to an image monage, find the tissue
	and return its centroid. Works so long as the tissue
	size is smaller than the resin.

	Returns supertile id
	"""
	while montage_image.ndim > 2:
		montage_image = montage_image[...,0]

	mips = tinybrain.downsample_with_averaging(montage_image, num_mips=1, factor=(2,2,1))
	montage_image = mips[0]

	labels = cc3d.connected_components(montage_image, delta=7)
	mask = cc3d.largest_k(labels, k=10) > 0
	mask = scipy.ndimage.binary_dilation(mask, iterations=10)

	roi = labels * (mask == 0)
	roi = roi > 0

	roi = cc3d.largest_k(roi, k=1)

	clipped_image = roi * montage_image
	second_labels = cc3d.connected_components(clipped_image, delta=7)
	holes = cc3d.largest_k(second_labels, k=5) > 0
	roi[holes] = 0

	if not np.any(roi):
		return None

	centroid = cc3d.statistics(roi)["centroids"][1]

	# dt = edt.edt(roi, parallel=True)
	# max_point = np.unravel_index(np.argmax(dt), dt.shape)
	# max_point = np.array(max_point)

	if visualize_roi:
		roi = roi.astype(np.uint8)
		centroid = centroid.astype(int)
		# max_point = max_point.astype(int)
		roi[centroid[0]-50:centroid[0]+50, centroid[1]-50:centroid[1]+50] = 2
		# roi[max_point[0]-50:max_point[0]+50, max_point[1]-50:max_point[1]+50] = 3
		microviewer.hyperview(montage_image, roi)

	# max_point *= 2
	# max_point //= np.array([ supertile_width, supertile_height ])

	centroid *= 2
	tc = np.copy(centroid)
	tc //= np.array([ supertile_width, supertile_height ])
	tc[0] += 1
	tc = tc.astype(int)

	return (
		supertile_map[tc[0], tc[1]],
		centroid
	)

def find_stage_zero(stage_info):
	zero = stage_info[0]

	for pos in stage_info:
		if pos["x_relroi_nm"] == 0 and pos["y_relroi_nm"] == 0:
			zero = pos
			break

	return np.array([zero["stage_x_nm"], zero["stage_y_nm"]], dtype=float)

def find_tissue_roi(section_path:str, visualize:bool = False) -> Tuple[Tuple[float,float], Tuple[float,float]]:
	"""
	Find a reasonable brightness/contrast normalization point
	and the tissue centroid given in TEM stage coordinates usable
	by voxa software.

	Where bcx, bcy means brightness/contrast x,y 
	and tissue means the whole tissue centroid

	Returns: 
		( (bcx, bcy), (tissue_x, tissue_y) )
	"""
	stage_info = read_stage_csv(section_path)
	montage_info = build_montage(section_path, stage_info)

	supertile_id, tissue_centroid = find_tissue_roi_kernel(
		montage_info["montage_image"],
		montage_info["supertile_width"],
		montage_info["supertile_height"],
		montage_info["supertile_map"],
		visualize
	)

	supertile_width_nm = abs(
		int(stage_info[0]["stage_x_nm"]) - int(stage_info[1]["stage_x_nm"])
	)

	null_return = (
		(None, None),
		(None, None)
	)

	if supertile_id is None:
		return null_return

	supertile = None
	for row in stage_info:
		if row["tile_id"] == supertile_id:
			supertile = row
			break

	supertile_coords = np.where(montage_info["supertile_map"][1:,:-1] == supertile_id)
	supertile_coords = np.array([ supertile_coords[0][0], supertile_coords[1][0] ])
	upper_corner = supertile_coords * montage_info["supertile_width"]
	
	tissue_centroid_supertile = tissue_centroid - upper_corner
	supertile_centroid = np.array([ montage_info["supertile_width"] // 2, montage_info["supertile_height"] // 2 ])

	delta = tissue_centroid_supertile - supertile_centroid

	supertile_centroid_nm = np.array([
		supertile["stage_x_nm"], supertile["stage_y_nm"]
	], dtype=float)

	nm_per_px = supertile_width_nm / montage_info["supertile_width"]
	delta_nm = delta * nm_per_px
	delta_nm[1] *= -1
	delta_nm = -delta_nm[::-1]

	stage_factor = 1e3

	true_centroid = supertile_centroid_nm + delta_nm

	return (
		supertile_centroid_nm / stage_factor,
		true_centroid / stage_factor
	)