Preparing training data from IDR using stardist

# Explore data
from IPython.display import display, HTML
import requests

IDR_BASE_URL = "https://idr.openmicroscopy.org"
MAP_URL = "{base}/webclient/api/annotations/?type=map&{type}={screen_id}"
PLATES_URL = "{base}/webclient/api/plates/?id={screen_id}"

# Here we need to set the screen ID for the dataset we want to explore. 
# The screen ID is a unique identifier for a specific dataset in the IDR (Image Data Resource). 
# You can find the screen ID in the URL of the dataset page on the IDR website. 
# For example, if the URL is "https://idr.openmicroscopy.org/webclient/?show=screen-1952", then the screen ID is 1952.
SCREEN_ID = 1952

INDEX_PAGE = f"{IDR_BASE_URL}/webclient/?experimenter=-1"
# create http session
with requests.Session() as session:
    request = requests.Request('GET', INDEX_PAGE)
    prepped = session.prepare_request(request)
    response = session.send(prepped)
    if response.status_code != 200:
        response.raise_for_status()
print(f"Connected to IDR at {IDR_BASE_URL} with status code {response.status_code}")

# Now we can print the metadata available on this screen
qs = {'base': IDR_BASE_URL, 'type': 'screen', 'screen_id': SCREEN_ID}
url = MAP_URL.format(**qs)
for a in session.get(url).json()['annotations']:
    namespace = a['ns']
    for v in a['values']:
        key = v[0]
        value = v[1]
        print(key, value)

qs = {'base': IDR_BASE_URL, 'screen_id': SCREEN_ID}
url = PLATES_URL.format(**qs)

for p in session.get(url).json()['plates']:
    plate_id = p['id']
    print("Plate ID: {id}, Name: {name}".format(**p))
    # create a dictionary with plated id and name

WELLS_IMAGES_URL = "{base}/webgateway/plate/{plate_id}/{field}/"
plate_id = 4403

qs = {'base': IDR_BASE_URL, 'plate_id': plate_id, 'field': 0}
url = WELLS_IMAGES_URL.format(**qs)
grid = session.get(url).json()
rowlabels = grid['rowlabels']
collabels = grid['collabels']
for r, row in enumerate(grid['grid']):
    for c, cell in enumerate(row):
        if cell is not None:
            well_name = f"{rowlabels[r]}{collabels[c]}"
            print(f"Well: {well_name}, Image: {cell['id']}")

# Now we can print the metadata available on a particular image
IMAGE_ID = 1898190
qs = {'base': IDR_BASE_URL, 'type': 'image', 'screen_id': IMAGE_ID}
url = MAP_URL.format(**qs)
for a in session.get(url).json()['annotations']:
    namespace = a['ns']
    for v in a['values']:
        key = v[0]
        value = v[1]
        print(key, value)

annotations = session.get(url).json()['annotations']

meta = {k: v for a in annotations for k, v in a["values"]}

channels = [c.split(":", 1) for c in meta["Channels"].split(";")]
channels = {i: (dye.strip(), target.strip()) for i, (dye, target) in enumerate(channels)}

for i, (dye, target) in channels.items():
    print(i, dye, "->", target)

# Pick one image

IMAGE_DETAILS_URL = "{base}/api/v0/m/images/{image_id}/"
IMAGE_ID = 1898190

qs = {'base': IDR_BASE_URL, 'image_id': IMAGE_ID}
url = IMAGE_DETAILS_URL.format(**qs)

r = session.get(url).json()

# Images backed by OME-Zarr will link to the zarr data via the image externalInfo
image_details = r["data"]["omero:details"]
is_zarr = "externalInfo" in image_details and image_details["externalInfo"]["EntityType"] == "com.glencoesoftware.ngff:multiscales"

print("Is Zarr?", is_zarr)
zarr_url = None
if (is_zarr):
    externalInfo_lsid = image_details["externalInfo"]["Lsid"]
    print(externalInfo_lsid)
    zarr_url = externalInfo_lsid.replace("s3:", "https:").replace("?anonymous=true", "")
    zarr_url = zarr_url.rstrip("/")
    display(HTML(f"Open <a href='https://ome.github.io/ome-ngff-validator/?source={zarr_url}'>image in NGFF-validator</a>"))

data = None
if zarr_url:
    import zarr
    import dask.array as da
    img_group = zarr.open(zarr_url)
    attrs = img_group.attrs.asdict()
    ds_paths = [ds["path"] for ds in attrs["multiscales"][0]["datasets"]]
    print("dataset paths", ds_paths)

    # Load dask array for largest dataset resolution
    data = da.from_zarr(f"{zarr_url}/{ds_paths[0]}")
    # T, C, Z, Y, X
    print("array shape", data.shape)

from matplotlib.pyplot import imshow
%matplotlib inline

channel = 3

if data is not None:
    # load the first plane (z, c, t = 0)
    extra_dims = [channel] * (data.ndim - 2)
    print("extra_dims", extra_dims)
    # we slice away any Z/C/T dims, leaving Y and X
    to_slice = (*extra_dims, slice(None), slice(None))
    plane = data[to_slice]
    print(plane.shape)
    # plane is a 2D numpy array
    imshow(plane)

#Get the plane and save as a numpy array
import numpy as np
import matplotlib.pyplot as plt

# data is T, C, Z, Y, X -> pick one channel, first timepoint and z-slice
channel = 2
img = np.asarray(plane)
print("plane shape", img.shape)

from stardist.models import StarDist2D
from csbdeep.utils import normalize

# Pretrained model for fluorescent nuclei, downloaded the first time it is used
model_nuclei = StarDist2D.from_pretrained('2D_versatile_fluo')

# scale > 1 enlarges the image before prediction, scale < 1 shrinks it
SCALE = 1.0

img_norm = normalize(img, 1, 99.8)
labels, _ = model_nuclei.predict_instances(img_norm, scale=SCALE)

print(f'found {labels.max()} nuclei')

from stardist import random_label_cmap

lbl_cmap = random_label_cmap()

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
axes[0].imshow(img, cmap='gray')
axes[0].set_title(f'channel {channel} (DAPI)')
axes[1].imshow(img, cmap='gray')
axes[1].imshow(labels, cmap=lbl_cmap, alpha=0.5)
axes[1].set_title(f'{labels.max()} nuclei')
for ax in axes:
    ax.axis('off')

from skimage.measure import regionprops

diameters = [p.equivalent_diameter_area for p in regionprops(labels)]
print(f'median nucleus diameter: {np.median(diameters):.1f} pixels')

from skimage.morphology import remove_small_objects
from skimage.segmentation import relabel_sequential

# Objects smaller than a quarter of the median nucleus area count as debris
min_area = 0.5 * np.median([p.area for p in regionprops(labels)])


def filter_small_objects(labels, min_area=min_area):
    """Remove objects of min_area pixels or smaller and renumber the labels 1..N."""
    filtered = remove_small_objects(labels, min_size=min_area)
    return relabel_sequential(filtered)[0]


labels_filtered = filter_small_objects(labels)
print(f'minimum area: {min_area:.0f} pixels')
print(f'{labels.max()} objects before, {labels_filtered.max()} after filtering')

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for ax, lbl, title in zip(axes, [labels, labels_filtered], ['before', 'after']):
    ax.imshow(img, cmap='gray')
    ax.imshow(lbl, cmap=lbl_cmap, alpha=0.5)
    ax.set_title(f'{title}: {lbl.max()} nuclei')
    ax.axis('off')

# Save the filtered labels for further analysis
np.save(f'labels_{IMAGE_ID}.npy', labels_filtered)
print('saved', f'labels_{IMAGE_ID}.npy')

import random

TRAIN_CHANNEL = 4    # phalloidin, the input for the network
NUCLEI_CHANNEL = 2   # DAPI, used only to make the labels
N_IMAGES = 100        # more images is better, but each one has to be downloaded

# We collect all image ids
image_ids = [c['id'] for row in grid['grid'] for c in row if c is not None]
print('images in plate:', len(image_ids))

random.seed(0)  # same subset every run

#from all images in the plate we 
subset = random.sample(image_ids, N_IMAGES)

# We create a function to load an image with a certain resolution from the IDR

def load_image(image_id, resolution=0):
    """Load an IDR image as a dask array (lowest resolution by default)."""
    url = IMAGE_DETAILS_URL.format(base=IDR_BASE_URL, image_id=image_id)
    lsid = session.get(url).json()['data']['omero:details']['externalInfo']['Lsid']
    zarr_url = lsid.replace('s3:', 'https:').replace('?anonymous=true', '').rstrip('/')
    paths = [ds['path'] for ds in zarr.open(zarr_url).attrs.asdict()['multiscales'][0]['datasets']]
    return da.from_zarr(f'{zarr_url}/{paths[resolution]}')

import zarr
import dask.array as da

# Download the subset once and keep it in memory (C, Y, X per image)
raw = [np.asarray(load_image(i)) for i in subset]
print('one image shape (C, Y, X):', raw[0].shape)

fig, axes = plt.subplots(2, 5, figsize=(15, 6))
for ax, im, image_id in zip(axes.ravel(), raw, subset):
    ax.imshow(im[TRAIN_CHANNEL], cmap='gray')
    ax.set_title(image_id, fontsize=8)
    ax.axis('off')
fig.suptitle(f'training input: channel {TRAIN_CHANNEL} (phalloidin)')

from pathlib import Path
import tifffile

OUT = Path('training_data')
(OUT / 'images').mkdir(parents=True, exist_ok=True)
(OUT / 'labels').mkdir(parents=True, exist_ok=True)
(OUT / 'nuclei').mkdir(parents=True, exist_ok=True)

def segment_nuclei(plane, scale=SCALE):
    """Normalize the DAPI plane and return the filtered StarDist nuclei as a label image."""
    labels, _ = model_nuclei.predict_instances(normalize(plane, 1, 99.8), scale=scale)
    labels = filter_small_objects(labels)
    return labels


def normalize_to_uint16(plane):
    """Percentile-normalize a plane and rescale it to fill the uint16 range.

    This is a cell-painting screen, so absolute DAPI brightness varies a lot between wells
    (imaging session, compound, exposure). BiaPy requires every raw image in a folder to fall
    in the same intensity "bucket" (e.g. all uint16-range, all 0-1 float): if one image happens
    to be dim enough that its raw max sits below 255, training crashes with a data-range error.
    Rescaling every image to use the full range removes that dependence on absolute brightness.
    """
    img_norm = np.clip(normalize(plane, 1, 99.8), 0, 1)
    return (img_norm * 65535).astype('uint16')


for image_id, im in zip(subset, raw):
    tifffile.imwrite(OUT / 'images' / f'{image_id}.tif', im[TRAIN_CHANNEL].astype('uint16'))
    tifffile.imwrite(OUT / 'labels' / f'{image_id}.tif', segment_nuclei(im[NUCLEI_CHANNEL]).astype('uint16'))
    # Raw DAPI intensity, kept as its own target for an image-to-image (virtual staining) model
    tifffile.imwrite(OUT / 'nuclei' / f'{image_id}.tif', normalize_to_uint16(im[NUCLEI_CHANNEL]))

print(f'wrote {len(subset)} image/label/nuclei triplets to {OUT.resolve()}')

import csv
import re

def image_metadata(image_id):
    """Plate, well and field from the image name; compound and dose from its annotations."""
    name = session.get(f'{IDR_BASE_URL}/webclient/imgData/{image_id}/').json()['meta']['imageName']
    plate, well, field = re.match(r'(\S+) \[Well (\w+), Field (\d+)\]', name).groups()
    url = MAP_URL.format(base=IDR_BASE_URL, type='image', screen_id=image_id)
    values = {k: v for a in session.get(url).json()['annotations'] for k, v in a['values']}
    return {'image_id': image_id, 'plate': plate, 'well': well, 'field': field,
            'compound': values.get('Compound Name', ''), 'dose': values.get('Dose', '')}

metadata = [image_metadata(image_id) for image_id in subset]

with open(OUT / 'metadata.csv', 'w', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=metadata[0].keys())
    writer.writeheader()
    writer.writerows(metadata)

print(f'wrote metadata for {len(metadata)} images to {OUT / "metadata.csv"}')