xenium_roi_extractor.py

"""
Module for extracting all data for a given list of rectangular ROIs
from Xenium Onboard Analysis output. The ROIs x,y,w,h are provided as part of
filenames in a foder (intended for images, but fake empty files can be used).
"""

import json
import os
import re
import argparse
from packaging.version import Version
import pandas as pd
import numpy as np
import tifffile
import h5py
import skimage
import scipy

def get_arguments():
 """
 Parses and checks command line arguments, and provides an help text.
 Assumes 3 and returns 3 positional command line arguments:
 """
 parser = argparse.ArgumentParser(
 description = "Split transcripts and images according to the given ROIs")
 parser.add_argument("xenium_folder", help = "path to the xenium folder")
 parser.add_argument("roi_folder", help = "path to the folder with the ROI images")
 parser.add_argument("output_folder", help = "path to the output folder")
 args = parser.parse_args()
 return args.xenium_folder, args.roi_folder, args.output_folder

def read_image_v2(xenium_dir):
 """
 Reads the morphology focus images from
 Xenium Onboard Analysis output v >= 2
 https://www.10xgenomics.com/support/software/xenium-onboard-analysis/latest/analysis/xoa-output-understanding-outputs#images
 """
 mf_tiff_file_path = os.path.join(xenium_dir, "morphology_focus",
 "morphology_focus_0000.ome.tif")
 print(f"Reading Morphology Focus Image File: {mf_tiff_file_path}")
 img = tifffile.imread(mf_tiff_file_path)
 return img

def process_image_v2(mf_image, roi, roi_name, output_dir):
 """
 Crops and saves the morphology focus images from
 Xenium Onboard Analysis output v > 2
 https://www.10xgenomics.com/support/software/xenium-onboard-analysis/latest/analysis/xoa-output-understanding-outputs
 """
 x, y, w, h = roi
 if mf_image.ndim== 3:
 sub_image = mf_image[:, y:y + h, x:x + w]
 for c, channel in enumerate(sub_image):
 outpath = os.path.join(output_dir, f"{roi_name}-morphology_focus_000{c}.ome.tif")
 tifffile.imwrite(outpath, data = channel, ome = True, compression = "lzw")
 else:
 sub_image = mf_image[y:y + h, x:x + w]
 outpath = os.path.join(output_dir, f"{roi_name}-morphology_focus_0000.ome.tif")
 tifffile.imwrite(outpath, data = sub_image, ome = True, compression = "lzw")

def read_image_v1(xenium_dir):
 """
 Reads the morphology mip image from
 Xenium Onboard Analysis output v <= 2
 https://www.10xgenomics.com/support/software/xenium-onboard-analysis/1.9/analysis/xoa-output-understanding-outputs
 """
 mip_tiff_file_path = os.path.join(xenium_dir, "morphology_mip.ome.tif")
 print(f"Reading Maximum Projection Image File: {mip_tiff_file_path}")
 img = tifffile.imread(mip_tiff_file_path)
 return img

def process_image_v1(mip_image, roi, roi_name, output_dir):
 """
 Crops and saves the morphology mip image from
 Xenium Onboard Analysis output v <= 2
 https://www.10xgenomics.com/support/software/xenium-onboard-analysis/1.9/analysis/xoa-output-understanding-outputs
 """
 x, y, w, h = roi
 sub_image = mip_image[:, y:y + h, x:x + w]
 tifffile.imwrite(os.path.join(output_dir, f"{roi_name}-morphology_mip.ome.tif"),
 data = sub_image, ome = True, compression = "lzw")

if __name__ == "__main__":

 xenium_folder, roi_folder, ouptut_folder = get_arguments()

 # Expects standard xenium filenames:
 # https://www.10xgenomics.com/support/in-situ-gene-expression/documentation/steps/onboard-analysis/at-a-glance-xenium-output-files
 xenium_file_path = os.path.join(xenium_folder, "experiment.xenium")
 # Reads the pixel size in um and the software version
 with open(xenium_file_path, encoding="utf-8") as metadata_file:
 metadata = json.load(metadata_file)
 pixel_size = metadata["pixel_size"]
 version = Version(metadata["instrument_sw_version"])

 transcript_file_path = os.path.join(xenium_folder, "transcripts.parquet")
 cells_file_path = os.path.join(xenium_folder, "cells.parquet")
 cell_boundaries_file_path = os.path.join(xenium_folder, "cell_boundaries.parquet")
 features_file_path = os.path.join(xenium_folder, "cell_feature_matrix.h5")
 print(f"Xenium file: {xenium_file_path}")
 print(f"Transcript file: {transcript_file_path}")
 print(f"Cell file: {cells_file_path}")
 print(f"Cell Boundary file: {cell_boundaries_file_path}")
 print(f"Feature Matrix file: {features_file_path}")

 # Reads the transcripts and add locations in pixels
 print(f"Reading Transcript file: {transcript_file_path}")
 transcripts = pd.read_parquet(transcript_file_path)
 transcripts["x_location_px"] = transcripts["x_location"] / pixel_size
 transcripts["y_location_px"] = transcripts["y_location"] / pixel_size
 transcripts["z_location_px"] = transcripts["z_location"] / pixel_size

 # Reads the cells and add locations in pixels
 print(f"Reading Cell file: {cells_file_path}")
 cells = pd.read_parquet(cells_file_path)
 cells["x_centroid_px"] = cells["x_centroid"] / pixel_size
 cells["y_centroid_px"] = cells["y_centroid"] / pixel_size

 # Reads the cells and add locations in pixels
 print(f"Reading Cell Boundary file: {cell_boundaries_file_path}")
 cell_boundaries = pd.read_parquet(cell_boundaries_file_path)
 cell_boundaries["vertex_x_px"] = cell_boundaries["vertex_x"] / pixel_size
 cell_boundaries["vertex_y_px"] = cell_boundaries["vertex_y"] / pixel_size

 # Reads ROIS, expects _Xnnn_Ynnnn_Wnnn_Hnnnn_ with n = 0-9
 # makes a lists of lists [x, y, w, h]
 print(f"Looking for ROIs in: {roi_folder}")
 ROIs = os.listdir(roi_folder)
 ROIs = filter(lambda filename: re.match(".*_X[0-9]+_Y[0-9]+_W[0-9]+_H[0-9]+.*", filename), ROIs)
 ROIs = [re.match(
 ".*_X([0-9]+)_Y([0-9]+)_W([0-9]+)_H([0-9]+).*", filename).groups() for filename in ROIs]
 ROIs = [list(int(n) for n in i) for i in ROIs]
 print(f"Found ROIs: {ROIs}")

 # Read features
 # Features are equivalent to:
 # transcripts.loc[(transcripts["qv"] > 20) & (transcripts["cell_id"] != "UNASSIGNED")]
 print(f"Processing features: {features_file_path}")
 with h5py.File(features_file_path, "r") as f:
 feature_attrs = dict(f.attrs.items())
 feature_h5 = {}
 feature_h5['barcodes'] = f["matrix"].get("barcodes")[:]
 feature_h5['data'] = f["matrix"].get("data")[:]
 feature_h5['indices'] = f["matrix"].get("indices")[:]
 feature_h5['indptr'] = f["matrix"].get("indptr")[:]
 feature_h5['shape'] = f["matrix"].get("shape")[:]
 feature_h5['features'] = {}
 feature_h5['features']["_all_tag_keys"] = \
 f["matrix"].get("features").get("_all_tag_keys")[:]
 feature_h5['features']["feature_type"] = f["matrix"].get("features").get("feature_type")[:]
 feature_h5['features']["genome"] = f["matrix"].get("features").get("genome")[:]
 feature_h5['features']["id"] = f["matrix"].get("features").get("id")[:]
 feature_h5['features']["name"] = f["matrix"].get("features").get("name")[:]
 features = scipy.sparse.csr_matrix(
 (feature_h5["data"], feature_h5["indices"], feature_h5["indptr"]), dtype=int)
 features = pd.DataFrame(features.toarray(), columns = feature_h5['features']["name"])
 features.columns = features.columns.astype(str)
 features["cell_id"] = cells["cell_id"]
 features = features[list(features.columns)[-1:] + list(features.columns)[:-1]]

 # Read the images
 if version < Version("2"):
 image = read_image_v1(xenium_folder)
 else:
 image = read_image_v2(xenium_folder)

 os.makedirs(ouptut_folder, exist_ok = True)

 for r in ROIs:
 ROI_STRING = "X_" + str(r[0]) + "_Y_" + str(r[1]) + "_W_" + str(r[2]) + "_H_"+ str(r[3])
 print(f"Processing: {ROI_STRING}")

 sub_transcripts = transcripts.loc[(transcripts['x_location_px'] >= r[0]) &
 (transcripts['x_location_px'] <= (r[0] + r[2])) &
 (transcripts['y_location_px'] >= r[1]) &
 (transcripts['y_location_px'] <= (r[1] + r[3]))]
 sub_transcripts.loc[:, "x_location_px_ROI"] = sub_transcripts["x_location_px"] - r[0]
 sub_transcripts.loc[:, "y_location_px_ROI"] = sub_transcripts["y_location_px"] - r[1]

 # We exlude all cells not fully contained in the ROI!
 cell_boundaries["out"] = (cell_boundaries['vertex_x_px'] >= r[0]) & \
 (cell_boundaries['vertex_x_px'] <= (r[0] + r[2])) & \
 (cell_boundaries['vertex_y_px'] >= r[1]) & \
 (cell_boundaries['vertex_y_px'] <= (r[1] + r[3]))
 filter_df = cell_boundaries[["cell_id", "out"]].groupby("cell_id").filter(
 lambda x: all(x["out"])).drop_duplicates()

 sub_cell_boundaries = pd.merge(cell_boundaries, filter_df["cell_id"],
 on = "cell_id", how = "inner")
 sub_cell_boundaries = sub_cell_boundaries.drop('out', axis=1)
 sub_cell_boundaries["vertex_x_px_ROI"] = sub_cell_boundaries["vertex_x_px"] - r[0]
 sub_cell_boundaries["vertex_y_px_ROI"] = sub_cell_boundaries["vertex_y_px"] - r[1]

 sub_cells = pd.merge(cells, filter_df["cell_id"], on = "cell_id", how = "inner")
 sub_cells["x_centroid_px_ROI"] = sub_cells["x_centroid_px"] - r[0]
 sub_cells["y_centroid_px_ROI"] = sub_cells["y_centroid_px"] - r[1]

 # Make a cell mask for the ROI by getting the coordinates of every pixel
 # inside each cell and then drawing them on ask of the size of the ROI
 # The cells are numbered in the same order as they appear in
 # the cell_boundaries.parquet file
 cell_coords = sub_cell_boundaries.groupby("cell_id").apply(
 lambda x, roi=r: skimage.draw.polygon(x["vertex_x_px_ROI"],
 x["vertex_y_px_ROI"], (roi[2], roi[3])))
 print(cell_coords)
 cell_mask = np.zeros((r[3], r[2]), dtype=np.int32)
 for n, i in enumerate(cell_coords):
 cell_mask[i[1], i[0]] = n

 # Process features
 sub_features = pd.merge(features, filter_df["cell_id"], on = "cell_id", how = "inner")
 fmatrix = sub_features.drop("cell_id", axis=1)
 fmatrix = scipy.sparse.csr_array(fmatrix)

 # Make transcript table for reding with the Polylux viewer (Resolve Biosciences)
 polylux_tx = sub_transcripts[["x_location_px_ROI", "y_location_px_ROI",
 "z_location_px", "feature_name", "qv"]]
 polylux_tx["x_location_px_ROI"] = polylux_tx["x_location_px_ROI"].astype(np.int32)
 polylux_tx["y_location_px_ROI"] = polylux_tx["y_location_px_ROI"].astype(np.int32)
 polylux_tx["z_location_px"] = polylux_tx["z_location_px"].astype(np.int32)

 # Output
 print(f"Writing {ROI_STRING}")
 subfolder= os.path.join(ouptut_folder, ROI_STRING)
 os.makedirs(subfolder, exist_ok = True)

 # Make h5 file
 with h5py.File(os.path.join(subfolder, ROI_STRING + "-cell_feature_matrix.h5"), "w") as f:
 for k,v in feature_attrs.items():
 f.attrs.create(k, v)
 m_to_w = f.create_group("matrix")
 m_to_w.create_dataset("barcodes", data=np.array(sub_features["cell_id"]))
 m_to_w.create_dataset("data", data=fmatrix.data)
 m_to_w.create_dataset("indices", data=fmatrix.indices)
 m_to_w.create_dataset("indptr", data=fmatrix.indptr)
 m_to_w.create_dataset("shape", data=fmatrix.shape)
 f_to_write = m_to_w.create_group("features")
 f_to_write.create_dataset("_all_tag_keys", data=feature_h5["features"]["_all_tag_keys"])
 f_to_write.create_dataset("feature_type", data=feature_h5["features"]["feature_type"])
 f_to_write.create_dataset("genome", data=feature_h5["features"]["genome"])
 f_to_write.create_dataset("id", data=feature_h5["features"]["id"])
 f_to_write.create_dataset("name", data=feature_h5["features"]["name"])

 if version < Version("2"):
 process_image_v1(image, r, ROI_STRING, subfolder)
 else:
 process_image_v2(image, r, ROI_STRING, subfolder)

 # Write other output
 sub_transcripts.to_parquet(os.path.join(subfolder, ROI_STRING + "-transcritps.parquet"))
 sub_transcripts.to_csv(os.path.join(subfolder, ROI_STRING + "-transcritps.csv"))
 polylux_tx.to_csv(os.path.join(subfolder, ROI_STRING + "-Polylux_transcritps.txt"),
 header = None, index = False, sep = "\t")
 sub_cells.to_parquet(os.path.join(subfolder, ROI_STRING + "-cells.parquet"))
 sub_cells.to_csv(os.path.join(subfolder, ROI_STRING + "-cells.csv"))
 sub_cell_boundaries.to_parquet(
 os.path.join(subfolder, ROI_STRING + "-cell_boundaries.parquet"))
 sub_cell_boundaries.to_csv(os.path.join(subfolder, ROI_STRING + "-cell_boundaries.csv"))
 sub_features.to_csv(os.path.join(subfolder, ROI_STRING + "-feature_matrix.csv"))
 tifffile.imwrite(os.path.join(subfolder, ROI_STRING + "-xenium_cell_mask.ome.tif"),
 data = cell_mask, ome = True, compression = "lzw")