Merge pull request #44 from rajewsky-lab/fix_imaging_preprocessing

Fix imaging preprocessing

docs/computational/generate_expression_matrix.md

@@ -13,7 +13,7 @@ To create such a spatial cell-by-gene ($M\times G$) expression matrix, you will
 
 We efficiently segment cells (or nuclei) from staining images using [cellpose](https:/MouseLand/cellpose).
 We provide a model that we fine-tuned for segmentation of fresh-frozen, H&E-stained tissue,
-[here](https://github.com/danilexn/openst/blob/main/models/HE_cellpose_rajewsky).
+[here](http://bimsbstatic.mdc-berlin.de/rajewsky/openst-public-data/models/HE_cellpose_rajewsky).
 You can specify any other model that works best for your data -
 refer to the [cellpose](https://cellpose.readthedocs.io/en/latest/index.html) documentation.
 
@@ -31,7 +31,7 @@ By default, segmentation is extended radially 10 pixels. This can be changed wit
 Make sure to populate the arguments with the values specific to your dataset. Here, we provide `--h5-in` consistent
 with the previous steps, `--image-in` and `--mask-out` will read and write the staining and mask inside the Open-ST h5 object,
 and `--model` is `HE_cellpose_rajewsky`, the default used in our manuscript. This is the model we recommend for H&E images, and
-weights are automatically downloaded. It is also [provided in our repo](https://github.com/rajewsky-lab/openst/blob/main/models/HE_cellpose_rajewsky).
+weights are automatically downloaded. It is also [provided in our repo](http://bimsbstatic.mdc-berlin.de/rajewsky/openst-public-data/models/HE_cellpose_rajewsky).
 The rest of parameters can be checked with `openst segment --help`.
 
 !!! tip

docs/computational/preprocessing_imaging.md

@@ -63,16 +63,33 @@ tile-scan. You can run this by running the following command on the stitched ima
 
 ```bash
 openst image_preprocess \
- --input=<path_to_input_image> \
- --CUT \
- --CUT-model=<path_to_model> \
- --output=<path_to_output>
+ --image-in Image_Stitched_Composite.tif \
+ --image-out Image_Stitched_Composite_Restored.tif
+ # --device cuda # in case you have a CUDA-compatible GPU
 ```
 
-Make sure to replace the placeholders (`<...>`). For instance,
-`<path_to_input_image>` is the full path and file name of the previously stitched image; `<path_to_model>`
-is filename our pre-trained [CUT model](https:/rajewsky-lab/openst/models/CUT.pth), and `<output_image>` 
-is the path to a folder (writeable) and desired filename for the output image.
+If you ran `openst merge_modalities`, then imaging data will be contained inside the Open-ST h5 object, and the command
+can be adapted:
+
+```bash
+openst image_preprocess \
+ --h5-in multimodal/spots_stitched.h5ad # just a placeholder, adapt
+ # --device cuda # in case you have a CUDA-compatible GPU
+```
+
+By default, the image will be loaded from the key `uns/spatial/staining_image`, and the CUT-restored image will be saved
+to `uns/spatial/staining_image_restored`. You can preview the image restoration results using:
+
+```bash
+openst preview \
+ --h5-in multimodal/spots_stitched.h5ad
+ --image-key uns/spatial/staining_image uns/spatial/staining_image_restored
+```
+
+This will load the two images and visualize it using `napari`. Later, you can run segmentation and pairwise alignment
+using either the default merged image (`uns/spatial/staining_image`), or the restored image (`uns/spatial/staining_image_restored`).
+Always assess these preprocessing choices (quantitatively and qualitatively) to decide whether these make sense for your data.
+
 
 ## Expected output
 After running the stitching (and optionally correction algorithms), you will have a single image file per sample. This, together with

docs/examples/adult_mouse/generate_expression_matrix.md

@@ -12,7 +12,7 @@ into a cell-by-genes matrix, where cells are defined from the segmentation mask.
 To create such a spatial cell-by-gene ($M\times G$) expression matrix, you will first need a segmentation mask.
 
 We efficiently segment cells (or nuclei) from staining images using [cellpose](https:/MouseLand/cellpose).
-For the H&E-stained tissue provided in this example, we used our [fine-tuned model](https://github.com/rajewsky-lab/openst/blob/main/models/HE_cellpose_rajewsky).
+For the H&E-stained tissue provided in this example, we used our [fine-tuned model](http://bimsbstatic.mdc-berlin.de/rajewsky/openst-public-data/models/HE_cellpose_rajewsky).
 Make sure to download it and save it into a new `models` folder that you need to create under the `openst_adult_demo` main folder.
 
 You can run the segmentation on the previously created `openst_demo_adult_mouse_spatial_beads_puck_collection_aligned.h5ad` file, which

docs/examples/e13_mouse/generate_expression_matrix.md

@@ -12,7 +12,7 @@ into a cell-by-genes matrix, where cells are defined from the segmentation mask.
 To create such a spatial cell-by-gene ($M\times G$) expression matrix, you will first need a segmentation mask.
 
 We efficiently segment cells (or nuclei) from staining images using [cellpose](https:/MouseLand/cellpose).
-For the H&E-stained tissue provided in this example, we used our [fine-tuned model](https://github.com/rajewsky-lab/openst/blob/main/models/HE_cellpose_rajewsky).
+For the H&E-stained tissue provided in this example, we used our [fine-tuned model](http://bimsbstatic.mdc-berlin.de/rajewsky/openst-public-data/models/HE_cellpose_rajewsky).
 Make sure to download it and save it into a new `models` folder that you need to create under the `openst_e13_demo` main folder.
 
 You can run the segmentation on the previously created `openst_demo_e13_mouse_head_spatial_beads_puck_collection_aligned.h5ad` file, which

openst/__init__.py

@@ -1 +1 @@
-__version__ = '0.0.2'
+__version__ = '0.0.11'

openst/cli.py

@@ -680,17 +680,46 @@ def get_image_preprocess_parser():
  allow_abbrev=False,
  add_help=False,
  )
- parser.add_argument("--input_img", type=str, required=True, help="path to input image")
- parser.add_argument("--cut_dir", type=str, required=True, help="path to CUT directory (to save patched images)")
- parser.add_argument("--tile_size_px", type=int, required=True, help="size of the tile in pixels")
+ parser.add_argument(
+ "--h5-in",
+ type=str,
+ default="",
+ help="""If set, image is loaded from the Open-ST h5 object (key in --image-in),
+ and retored image is saved there (to the key --image-out)""",
+ )
+ parser.add_argument(
+ "--image-in",
+ type=str,
+ default="uns/spatial/staining_image",
+ help="Key or path to the input image",
+ )
+ parser.add_argument(
+ "--image-out",
+ type=str,
+ default="uns/spatial/staining_image_restored",
+ help="Key or path where the restored image will be written into",
+ )
+ parser.add_argument(
+ "--tile-size-px",
+ type=int,
+ default=512,
+ help="The input image is split into squared tiles of side `--tile-size-px`, for inference."+ 
+ "Larger values avoid boundary effects, but require more memory.",
+ )
+ parser.add_argument("--model", type=str, default="HE_CUT_rajewsky", help="CUT model used for image restoration")
  parser.add_argument(
  "--device",
  type=str,
  default="cpu",
  choices=["cpu", "cuda"],
- help="Device used to run feature matching model. Can be ['cpu', 'cuda']",
+ help="Device used to run CUT restoration model. Can be ['cpu', 'cuda']",
+ )
+ parser.add_argument(
+ "--num-workers",
+ type=int,
+ default=-1,
+ help="Number of CPU workers for parallel processing",
  )
- parser.add_argument("--checkpoints_dir", type=str, default="./checkpoints", help="models are saved here")
 
  return parser
 

openst/preprocessing/CUT/models/__init__.py

@@ -19,8 +19,9 @@
 """
 
 import importlib
-from models.base_model import BaseModel
+import logging
 
+from openst.preprocessing.CUT.models.base_model import BaseModel
 
 def find_model_using_name(model_name):
  """Import the module "models/[model_name]_model.py".
@@ -29,7 +30,7 @@ def find_model_using_name(model_name):
  be instantiated. It has to be a subclass of BaseModel,
  and it is case-insensitive.
  """
- model_filename = "models." + model_name + "_model"
+ model_filename = "openst.preprocessing.CUT.models." + model_name + "_model"
  modellib = importlib.import_module(model_filename)
  model = None
  target_model_name = model_name.replace('_', '') + 'model'
@@ -63,5 +64,5 @@ def create_model(opt):
  """
  model = find_model_using_name(opt.model)
  instance = model(opt)
- print("model [%s] was created" % type(instance).__name__)
+ logging.info("Model architecture `%s` was created" % type(instance).__name__)
  return instance

openst/preprocessing/CUT/models/base_model.py

@@ -1,8 +1,10 @@
 import os
 import torch
+import logging
 from collections import OrderedDict
 from abc import ABC, abstractmethod
-from . import networks
+
+from openst.preprocessing.CUT.models import networks
 
 
 class BaseModel(ABC):
@@ -211,7 +213,7 @@ def load_networks(self, epoch):
  net = getattr(self, 'net' + name)
  if isinstance(net, torch.nn.DataParallel):
  net = net.module
- print('loading the model from %s' % load_path)
+ logging.info(f'Loading model weights from {load_path}')
  # if you are using PyTorch newer than 0.4 (e.g., built from
  # GitHub source), you can remove str() on self.device
  state_dict = torch.load(load_path, map_location=str(self.device))
@@ -229,17 +231,18 @@ def print_networks(self, verbose):
  Parameters:
  verbose (bool) -- if verbose: print the network architecture
  """
- print('---------- Networks initialized -------------')
+ message = '---------- (Start) Networks initialized -------------\n'
  for name in self.model_names:
  if isinstance(name, str):
  net = getattr(self, 'net' + name)
  num_params = 0
  for param in net.parameters():
  num_params += param.numel()
  if verbose:
- print(net)
- print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
- print('-----------------------------------------------')
+ message += f"{net}\n"
+ message += '[Network %s] Total number of parameters : %.3f M\n' % (name, num_params / 1e6)
+ message += '---------- (End) Networks initialized -------------'
+ logging.debug(message)
 
  def set_requires_grad(self, nets, requires_grad=False):
  """Set requies_grad=Fasle for all the networks to avoid unnecessary computations

openst/preprocessing/CUT/models/cut_model.py

@@ -1,10 +1,10 @@
 import numpy as np
 import torch
-from .base_model import BaseModel
-from . import networks
-from .patchnce import PatchNCELoss
-import util.util as util
 
+from openst.preprocessing.CUT.models.base_model import BaseModel
+from openst.preprocessing.CUT.models import networks
+from openst.preprocessing.CUT.models.patchnce import PatchNCELoss
+from openst.preprocessing.CUT.util import util
 
 class CUTModel(BaseModel):
  """ This class implements CUT and FastCUT model, described in the paper
@@ -98,10 +98,10 @@ def data_dependent_initialize(self, data):
  initialized at the first feedforward pass with some input images.
  Please also see PatchSampleF.create_mlp(), which is called at the first forward() call.
  """
- bs_per_gpu = data["A"].size(0) // max(len(self.opt.gpu_ids), 1)
+ bs_per_gpu = self.opt.batch_size
  self.set_input(data)
  self.real_A = self.real_A[:bs_per_gpu]
- self.real_B = self.real_B[:bs_per_gpu]
+ self.real_B = None
  self.forward() # compute fake images: G(A)
  if self.opt.isTrain:
  self.compute_D_loss().backward() # calculate gradients for D
@@ -138,14 +138,11 @@ def set_input(self, input):
  input (dict): include the data itself and its metadata information.
  The option 'direction' can be used to swap domain A and domain B.
  """
- AtoB = self.opt.direction == 'AtoB'
- self.real_A = input['A' if AtoB else 'B'].to(self.device)
- self.real_B = input['B' if AtoB else 'A'].to(self.device)
- self.image_paths = input['A_paths' if AtoB else 'B_paths']
+ self.real_A = input['A'].to(self.device)
 
  def forward(self):
  """Run forward pass; called by both functions <optimize_parameters> and <test>."""
- self.real = torch.cat((self.real_A, self.real_B), dim=0) if self.opt.nce_idt and self.opt.isTrain else self.real_A
+ self.real = self.real_A
  if self.opt.flip_equivariance:
  self.flipped_for_equivariance = self.opt.isTrain and (np.random.random() < 0.5)
  if self.flipped_for_equivariance:

openst/preprocessing/CUT/models/networks.py

@@ -5,12 +5,6 @@
 import functools
 from torch.optim import lr_scheduler
 import numpy as np
-from .stylegan_networks import StyleGAN2Discriminator, StyleGAN2Generator, TileStyleGAN2Discriminator
-
-###############################################################################
-# Helper Functions
-###############################################################################
-
 
 def get_filter(filt_size=3):
  if(filt_size == 1):
@@ -256,10 +250,6 @@ def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, in
  net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
  elif netG == 'unet_256':
  net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
- elif netG == 'stylegan2':
- net = StyleGAN2Generator(input_nc, output_nc, ngf, use_dropout=use_dropout, opt=opt)
- elif netG == 'smallstylegan2':
- net = StyleGAN2Generator(input_nc, output_nc, ngf, use_dropout=use_dropout, n_blocks=2, opt=opt)
  elif netG == 'resnet_cat':
  n_blocks = 8
  net = G_Resnet(input_nc, output_nc, opt.nz, num_downs=2, n_res=n_blocks - 4, ngf=ngf, norm='inst', nl_layer='relu')
@@ -323,8 +313,6 @@ def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal'
  net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, no_antialias=no_antialias,)
  elif netD == 'pixel': # classify if each pixel is real or fake
  net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
- elif 'stylegan2' in netD:
- net = StyleGAN2Discriminator(input_nc, ndf, n_layers_D, no_antialias=no_antialias, opt=opt)
  else:
  raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
  return init_net(net, init_type, init_gain, gpu_ids,
@@ -1214,8 +1202,6 @@ def forward(self, input):
 
 class UnetSkipConnectionBlock(nn.Module):
  """Defines the Unet submodule with skip connection.
- X -------------------identity----------------------
- |-- downsampling -- |submodule| -- upsampling --|
  """
 
  def __init__(self, outer_nc, inner_nc, input_nc=None,

openst/preprocessing/CUT/models/patchnce.py

@@ -0,0 +1,54 @@
+from packaging import version
+import torch
+from torch import nn
+
+class PatchNCELoss(nn.Module):
+ def __init__(self, opt):
+ super().__init__()
+ self.opt = opt
+ self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none')
+ self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool
+
+ def forward(self, feat_q, feat_k):
+ num_patches = feat_q.shape[0]
+ dim = feat_q.shape[1]
+ feat_k = feat_k.detach()
+
+ # pos logit
+ l_pos = torch.bmm(
+ feat_q.view(num_patches, 1, -1), feat_k.view(num_patches, -1, 1))
+ l_pos = l_pos.view(num_patches, 1)
+
+ # neg logit
+
+ # Should the negatives from the other samples of a minibatch be utilized?
+ # In CUT and FastCUT, we found that it's best to only include negatives
+ # from the same image. Therefore, we set
+ # --nce_includes_all_negatives_from_minibatch as False
+ # However, for single-image translation, the minibatch consists of
+ # crops from the "same" high-resolution image.
+ # Therefore, we will include the negatives from the entire minibatch.
+ if self.opt.nce_includes_all_negatives_from_minibatch:
+ # reshape features as if they are all negatives of minibatch of size 1.
+ batch_dim_for_bmm = 1
+ else:
+ batch_dim_for_bmm = self.opt.batch_size
+
+ # reshape features to batch size
+ feat_q = feat_q.view(batch_dim_for_bmm, -1, dim)
+ feat_k = feat_k.view(batch_dim_for_bmm, -1, dim)
+ npatches = feat_q.size(1)
+ l_neg_curbatch = torch.bmm(feat_q, feat_k.transpose(2, 1))
+
+ # diagonal entries are similarity between same features, and hence meaningless.
+ # just fill the diagonal with very small number, which is exp(-10) and almost zero
+ diagonal = torch.eye(npatches, device=feat_q.device, dtype=self.mask_dtype)[None, :, :]
+ l_neg_curbatch.masked_fill_(diagonal, -10.0)
+ l_neg = l_neg_curbatch.view(-1, npatches)
+
+ out = torch.cat((l_pos, l_neg), dim=1) / self.opt.nce_T
+
+ loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long,
+ device=feat_q.device))
+
+ return loss

openst/preprocessing/CUT/models/template_model.py

@@ -9,15 +9,16 @@
 Given input-output pairs (data_A, data_B), it learns a network netG that can minimize the following L1 loss:
  min_<netG> ||netG(data_A) - data_B||_1
 You need to implement the following functions:
- <modify_commandline_options>:　Add model-specific options and rewrite default values for existing options.
+ <modify_commandline_options>: Add model-specific options and rewrite default values for existing options.
  <__init__>: Initialize this model class.
  <set_input>: Unpack input data and perform data pre-processing.
  <forward>: Run forward pass. This will be called by both <optimize_parameters> and <test>.
  <optimize_parameters>: Update network weights; it will be called in every training iteration.
 """
 import torch
-from .base_model import BaseModel
-from . import networks
+
+from openst.preprocessing.CUT.models.base_model import BaseModel
+from openst.preprocessing.CUT.models import networks
 
 
 class TemplateModel(BaseModel):