Upload 26 files

Dockerfile

@@ -0,0 +1,27 @@
+FROM python:3.9 as build
+
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r ./requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
+# RUN mim install mmengine
+# RUN mim install "mmcv==2.1.0" & mim install "mmdet==3.3.0"
+
+FROM build as final
+
+COPY  --chown=user . .
+
+CMD python app.py

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,32 @@
 ---
 title: GRANA
-emoji: 🐨
-colorFrom: indigo
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.4.0
 app_file: app.py
-pinned: false
-short_description: Graphical Recognition&Analysis of Nanostructural Assemblies
+sdk: gradio
+sdk_version: 4.44.0
 ---
+# GRANA
+<img src="https://img.shields.io/badge/Python-3.9-blue"/>
+<a href="www.chloroplast.pl/GRANA"><img src="https://img.shields.io/badge/GRANA-Website-green" /></a>
+<a href="https://huggingface.co/spaces/chloroplast/GRANA"><img src="https://img.shields.io/badge/GRANA-Demo-green" /></a>
+<img src="https://img.shields.io/badge/Gradio-4.44.0-darkgreen"/>
+
+GRANA (**G**raphical **R**ecognition and **A**nalysis of **N**anostructural **A**ssemblies) 
+is an an AI-enhanced, user-friendly
+software tool that recognizes grana on thylakoid network electron micrographs 
+and generates a complex set of their structural parameters measurements.
+
+## Website
+More information about GRANA, including **example dataset**, can be found at [GRANA website](https://www.chloroplast.pl/grana).
+
+## Demo
+Demo version of GRANA is available at [Hugging Face Spaces](https://huggingface.co/spaces/chloroplast/GRANA).
+Using demo version, you can analyze up to 5 images at once.
+
+## Running as Docker container
+The recommended way to run GRANA is to use Docker container.
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+To run the container, use the following command:
+```bash
+docker run -p 7860:7860 mbuk/grana_measure:v0.5.4
+```
+After running the command, you can access the GRANA interface at `http://localhost:7860`.

angle_calculation/angle_model.py

@@ -0,0 +1,444 @@
+import numpy as np
+import torch
+import pytorch_lightning as pl
+import timm
+
+# from hydra.utils import instantiate
+from scipy.stats import circmean, circstd
+from scipy import ndimage
+from skimage.transform import resize
+
+from sampling import get_crop_batch
+from granum_utils import get_circle_mask
+import image_transforms
+from envelope_correction import calculate_best_angle_from_mask
+## loss
+
+class ConfidenceScaler:
+    def __init__(self, data: np.ndarray):
+        self.data = data
+        self.data.sort()
+    def __call__(self, x):
+        return np.searchsorted(self.data,x) / len(self.data)
+
+class PatchedPredictor:
+    def __init__(self,
+                 model,
+                 crop_size=96,
+                 normalization=dict(mean=0,std=1),
+                 n_samples=32,
+                 mask=None,# 'circle', None
+                 filter_outliers=True,
+                 apply_radon=False, # apply Radon transform
+                 radon_size=(128,128), # (int, int) reshape radon transformed image to this shape,
+                 angle_confidence_threshold=0,
+                 use_envelope_correction=True
+                ):
+        self.model = model
+        self.crop_size = crop_size
+        self.normalization = normalization
+        self.n_samples = n_samples
+        if mask not in [None, 'circle']:
+            raise ValueError(f'unknown mask {mask}')
+        self.mask = mask
+        self.filter_outliers = filter_outliers
+        
+        self.apply_radon = apply_radon
+        self.radon_size = radon_size
+        
+        self.angle_confidence_threshold = angle_confidence_threshold
+        self.use_envelope_correction = use_envelope_correction
+        
+    @torch.no_grad()
+    def __call__(self, img: np.ndarray, mask: np.ndarray):
+        pl.seed_everything(44)
+        # get crops with different scales and rotation
+        crops, angles_tta, scales_tta = get_crop_batch(
+            img, mask,
+            crop_size=self.crop_size,
+            samples_per_scale=self.n_samples,
+            use_variance_threshold=True
+        )
+        if len(crops) == 0:
+            return dict(
+                est_angle=np.nan,
+                est_angle_confidence=0.,
+            )
+            
+        # preprocess batch (normalize, mask, transform)
+        batch = self._preprocess_batch(crops)           
+        
+        # predict for batch - we don't use period and lumen anymore
+        preds_direction, preds_period, preds_lumen_width = self.model(batch)
+        # # convert to numpy
+        # preds_direction = preds_direction.numpy()
+        # preds_period = preds_period.numpy()
+        # preds_lumen_width = preds_lumen_width.numpy()
+        
+        # aggregate angles
+        est_angles = (preds_direction - angles_tta) % 180
+        est_angle = circmean(est_angles, low=-90, high=90) + 90
+        est_angle_std = circstd(est_angles, low=-90, high=90)
+        est_angle_confidence = self._std_to_confidence(est_angle_std, 10) # confidence 0.5 for std =10 degrees        
+        
+        if est_angle_confidence < self.angle_confidence_threshold:
+            est_angle = np.nan
+            est_angle_confidence = 0.
+        
+        if self.use_envelope_correction and (not np.isnan(est_angle)):           
+            angle_correction = -calculate_best_angle_from_mask(
+                ndimage.rotate(mask, -est_angle, reshape=True, order=0)
+            )            
+            est_angle += angle_correction
+            
+        return dict(
+            est_angle=est_angle,
+            est_angle_confidence=est_angle_confidence,
+        )
+        
+    def _apply_radon(self, batch): # may reauire circle mask
+        crops_radon = image_transforms.batched_radon(batch.numpy())
+        crops_radon = np.transpose(resize(np.transpose(crops_radon, (1, 2, 0)), self.radon_size), (2, 0, 1))
+        return torch.tensor(crops_radon)
+    
+    def _preprocess_batch(self, batch):
+        if self.mask == 'circle':
+            mask = get_circle_mask(batch.shape[1])
+            batch[:,mask] = 0
+        if self.apply_radon:
+            batch = self._apply_radon(batch)
+        batch = ((batch/255) - self.normalization['mean'])/self.normalization['std']
+        return batch.unsqueeze(1) # add channel dimension
+
+    def _filter_outliers(self, x, qmin=0.25, qmax=0.75):
+        x_min, x_max = np.quantile(x, [qmin, qmax])
+        return x[(x>=x_min) & (x<=x_max)]
+    
+    def _std_to_confidence(self, x, x_thr, y_thr=0.5):
+        """transform [0, inf] to [1,0], such that f(x_thr)=y_thr"""
+        return 1 / (1+x*(1-y_thr)/(x_thr*y_thr))
+
+class CosineLoss(torch.nn.Module):
+  def __init__(self, p=1, degrees=False, scale=1):
+    super().__init__()
+    self.p = p
+    self.degrees = degrees
+    self.scale = scale
+  def forward(self, x, y):
+    if self.degrees:
+      x = torch.deg2rad(x)
+      y = torch.deg2rad(y)
+    return torch.mean((1-torch.cos(x-y))**self.p) * self.scale
+    
+## model
+class AngleParser2d(torch.nn.Module):
+    def __init__(self, angle_range=180):
+        super().__init__()
+        self.angle_range = angle_range
+    def forward(self, batch):
+        # r = torch.linalg.norm(batch, dim=1)
+        preds_y_proj = torch.sigmoid(batch[:,0]) - 0.5
+        preds_x_proj = torch.sigmoid(batch[:,1]) - 0.5
+        preds_direction = self.angle_range/360.*torch.rad2deg(torch.arctan2(preds_y_proj, preds_x_proj))
+        return preds_direction
+
+class AngleRegularizer(torch.nn.Module):
+  def __init__(self, strength=1.0, scale=1.0, p=2):
+    super().__init__()
+    self.strength = strength
+    self.scale = scale
+    self.p = p
+  def forward(self, batch):
+    r = torch.linalg.norm(batch, dim=1)
+    return self.strength * torch.norm(r - self.scale, p=self.p)
+
+class AngleRegularizerLog(torch.nn.Module):
+  def __init__(self, strength=1.0, scale=1.0, p=2):
+    super().__init__()
+    self.strength = strength
+    self.scale = scale
+    self.p = p
+  def forward(self, batch):
+    r = torch.linalg.norm(batch, dim=1)
+    return self.strength * torch.norm(torch.log(r/self.scale), p=self.p)
+
+class StripsModel(pl.LightningModule):
+  def __init__(self, 
+              model_name = 'resnet18',
+               lr=0.001,
+               optimizer_hparams=dict(),
+               lr_hparams=dict(classname='MultiStepLR', kwargs=dict(milestones=[100, 150], gamma=0.1)),
+               loss_hparams=dict(rotation_weight=10., lumen_fraction_weight=50.),
+               angle_hparams=dict(angle_range=180.),
+               regularizer_hparams=None,
+               sigmoid_smoother=10.
+               ):
+    super().__init__()
+    # Exports the hyperparameters to a YAML file, and create "self.hparams" namespace
+    self.save_hyperparameters()
+    # Create model - implemented in non-abstract classes
+    self.model =  timm.create_model(model_name, in_chans=1, num_classes=4) #2 + self.hparams.angle_hparams['ndim'])
+    self.angle_parser = AngleParser2d(**self.hparams.angle_hparams)
+    self.regularizer = self._get_regularizer(self.hparams.regularizer_hparams)
+    self.losses = {
+        'direction': CosineLoss(2., True),
+        'period': torch.nn.functional.mse_loss,
+        'lumen_fraction': torch.nn.functional.mse_loss
+    }
+    self.losses_weights = {
+      'direction': self.hparams.loss_hparams['rotation_weight'],
+      'period': 1,
+      'lumen_fraction': self.hparams.loss_hparams['lumen_fraction_weight'],
+      'regularization': self.hparams.loss_hparams.get('regularization_weight', 0.)
+    }
+  
+  def _get_regularizer(self, regularizer_params):
+    if regularizer_params is None:
+      return None
+    else:
+      return instantiate(regularizer_params)
+    
+
+  def forward(self, x, return_raw=False):
+    """get predictions from image batch"""
+    preds = self.model(x) # preds: logit angle_sin, logit angle_cos, period, logit lumen fraction or logit angle, period, logit lumen fraction
+    preds_direction = self.angle_parser(preds)
+    preds_period = preds[:,-2]
+    preds_lumen_fraction = torch.sigmoid(preds[:,-1]*self.hparams.sigmoid_smoother) #lumen fraction is between 0 and 1, so we take sigmoid fo this
+
+    outputs = [preds_direction, preds_period, preds_lumen_fraction]
+    if return_raw:
+      outputs.append(preds)
+      
+    return tuple(outputs)
+
+  def configure_optimizers(self):
+    # AdamW is Adam with a correct implementation of weight decay (see here
+    # for details: https://arxiv.org/pdf/1711.05101.pdf)    
+    optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr, **self.hparams.optimizer_hparams)
+    # scheduler = getattr(torch.optim.lr_scheduler, self.hparams.lr_hparams['classname'])(optimizer, **self.hparams.lr_hparams['kwargs'])
+    scheduler = instantiate({**self.hparams.lr_hparams, '_partial_': True})(optimizer)
+    return [optimizer], [scheduler]
+
+  def process_batch_supervised(self, batch):
+    """get predictions, losses and mean errors (MAE)"""
+
+    # get predictions
+    preds = {}
+    preds['direction'], preds['period'], preds['lumen_fraction'], preds_raw = self.forward(batch['image'], return_raw=True) # preds: angle, period, lumen fraction, raw preds
+
+    # calculate losses
+    losses = {
+        'direction': self.losses['direction'](2*batch['direction'], 2*preds['direction']),
+        'period': self.losses['period'](batch['period'], preds['period']),
+        'lumen_fraction': self.losses['lumen_fraction'](batch['lumen_fraction'], preds['lumen_fraction']),
+    }
+    if self.regularizer is not None:
+      losses['regularization'] = self.regularizer(preds_raw[:,:2])
+    
+    losses['final'] = \
+      losses['direction']*self.losses_weights['direction'] + \
+      losses['period']*self.losses_weights['period'] + \
+      losses['lumen_fraction']*self.losses_weights['lumen_fraction'] + \
+      losses.get('regularization', 0.)*self.losses_weights.get('regularization', 0.)
+
+    # calculate mean errors
+    period_difference = np.mean(abs(
+      batch['period'].detach().cpu().numpy() - \
+      preds['period'].detach().cpu().numpy()
+    ))
+
+    a1 = batch['direction'].detach().cpu().numpy()
+    a2 = preds['direction'].detach().cpu().numpy()
+    angle_difference = np.mean(0.5*np.degrees(np.arccos(np.cos(2*np.radians(a2-a1)))))
+
+    lumen_fraction_difference = np.mean(abs(preds['lumen_fraction'].detach().cpu().numpy()-batch['lumen_fraction'].detach().cpu().numpy()))
+
+    mae = {
+      'period': period_difference,
+      'direction': angle_difference,
+      'lumen_fraction': lumen_fraction_difference
+    }
+
+    return preds, losses, mae
+
+  def log_all(self, losses, mae, prefix=''):
+    self.log(f"{prefix}angle_loss", losses['direction'].item())
+    self.log(f"{prefix}period_loss", losses['period'].item())
+    self.log(f"{prefix}lumen_fraction_loss", losses['lumen_fraction'].item())
+    self.log(f"{prefix}period_difference", mae['period'])
+    self.log(f"{prefix}angle_difference", mae['direction'])
+    self.log(f"{prefix}lumen_fraction_difference", mae['lumen_fraction'])
+    self.log(f"{prefix}loss", losses['final'])
+    if 'regularization' in losses:
+      self.log(f"{prefix}regularization_loss", losses['regularization'].item())
+  
+  def training_step(self, batch, batch_idx):
+    # "batch" is the output of the training data loader.
+    preds, losses, mae = self.process_batch_supervised(batch)
+    self.log_all(losses, mae, prefix='train_')
+
+    return losses['final'] 
+  
+  def validation_step(self, batch, batch_idx):
+    preds, losses, mae = self.process_batch_supervised(batch)
+    self.log_all(losses, mae, prefix='val_')
+  
+  def test_step(self, batch, batch_idx):
+    preds, losses, mae = self.process_batch_supervised(batch)
+    self.log_all(losses, mae, prefix='test_')
+
+    
+class StripsModelLumenWidth(pl.LightningModule):
+  def __init__(self, 
+              model_name = 'resnet18',
+               lr=0.001,
+               optimizer_hparams=dict(),
+               lr_hparams=dict(classname='MultiStepLR', kwargs=dict(milestones=[100, 150], gamma=0.1)),
+               loss_hparams=dict(rotation_weight=10., lumen_width_weight=50.),
+               angle_hparams=dict(angle_range=180.),
+               regularizer_hparams=None,
+               sigmoid_smoother=10.
+               ):
+    super().__init__()
+    # Exports the hyperparameters to a YAML file, and create "self.hparams" namespace
+    self.save_hyperparameters()
+    # Create model - implemented in non-abstract classes
+    self.model =  timm.create_model(model_name, in_chans=1, num_classes=4) #2 + self.hparams.angle_hparams['ndim'])
+    self.angle_parser = AngleParser2d(**self.hparams.angle_hparams)
+    self.regularizer = self._get_regularizer(self.hparams.regularizer_hparams)
+    self.losses = {
+        'direction': CosineLoss(2., True),
+        'period': torch.nn.functional.mse_loss,
+        'lumen_width': torch.nn.functional.mse_loss
+    }
+    self.losses_weights = {
+      'direction': self.hparams.loss_hparams['rotation_weight'],
+      'period': 1,
+      'lumen_width': self.hparams.loss_hparams['lumen_width_weight'],
+      'regularization': self.hparams.loss_hparams.get('regularization_weight', 0.)
+    }
+  
+  def _get_regularizer(self, regularizer_params):
+    if regularizer_params is None:
+      return None
+    else:
+      return instantiate(regularizer_params)
+
+  def forward(self, x, return_raw=False):
+    """get predictions from image batch"""
+    preds = self.model(x) # preds: logit angle_sin, logit angle_cos, period, logit lumen fraction or logit angle, period, logit lumen fraction
+    preds_direction = self.angle_parser(preds)
+    preds_period = preds[:,-2]
+    preds_lumen_width = preds[:,-1] #lumen fraction is between 0 and 1, so we take sigmoid fo this
+
+    outputs = [preds_direction, preds_period, preds_lumen_width]
+    if return_raw:
+      outputs.append(preds)
+      
+    return tuple(outputs)
+
+  def configure_optimizers(self):
+    # AdamW is Adam with a correct implementation of weight decay (see here
+    # for details: https://arxiv.org/pdf/1711.05101.pdf)    
+    optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr, **self.hparams.optimizer_hparams)
+    # scheduler = getattr(torch.optim.lr_scheduler, self.hparams.lr_hparams['classname'])(optimizer, **self.hparams.lr_hparams['kwargs'])
+    scheduler = instantiate({**self.hparams.lr_hparams, '_partial_': True})(optimizer)
+    return [optimizer], [scheduler]
+
+  def process_batch_supervised(self, batch):
+    """get predictions, losses and mean errors (MAE)"""
+
+    # get predictions
+    preds = {}
+    preds['direction'], preds['period'], preds['lumen_width'], preds_raw = self.forward(batch['image'], return_raw=True) # preds: angle, period, lumen fraction, raw preds
+
+    # calculate losses
+    losses = {
+        'direction': self.losses['direction'](2*batch['direction'], 2*preds['direction']),
+        'period': self.losses['period'](batch['period'], preds['period']),
+        'lumen_width': self.losses['lumen_width'](batch['lumen_width'], preds['lumen_width']),
+    }
+    if self.regularizer is not None:
+      losses['regularization'] = self.regularizer(preds_raw[:,:2])
+    
+    losses['final'] = \
+      losses['direction']*self.losses_weights['direction'] + \
+      losses['period']*self.losses_weights['period'] + \
+      losses['lumen_width']*self.losses_weights['lumen_width'] + \
+      losses.get('regularization', 0.)*self.losses_weights.get('regularization', 0.)
+
+    # calculate mean errors
+    period_difference = np.mean(abs(
+      batch['period'].detach().cpu().numpy() - \
+      preds['period'].detach().cpu().numpy()
+    ))
+
+    a1 = batch['direction'].detach().cpu().numpy()
+    a2 = preds['direction'].detach().cpu().numpy()
+    angle_difference = np.mean(0.5*np.degrees(np.arccos(np.cos(2*np.radians(a2-a1)))))
+
+    lumen_width_difference = np.mean(abs(preds['lumen_width'].detach().cpu().numpy()-batch['lumen_width'].detach().cpu().numpy()))
+    
+    lumen_fraction_pred = preds['lumen_width'].detach().cpu().numpy()/preds['period'].detach().cpu().numpy()
+    lumen_fraction_gt = batch['lumen_width'].detach().cpu().numpy()/batch['period'].detach().cpu().numpy()
+    lumen_fraction_difference = np.mean(abs(lumen_fraction_pred-lumen_fraction_gt))
+
+    mae = {
+      'period': period_difference,
+      'direction': angle_difference,
+      'lumen_width': lumen_width_difference,
+      'lumen_fraction': lumen_fraction_difference
+    }
+
+    return preds, losses, mae
+
+  def log_all(self, losses, mae, prefix=''):
+    for k, v in losses.items():
+        self.log(f'{prefix}{k}_loss', v.item() if isinstance(v, torch.Tensor) else v)
+    for k, v in mae.items():
+        self.log(f'{prefix}{k}_difference', v.item() if isinstance(v, torch.Tensor) else v)
+  
+  def training_step(self, batch, batch_idx):
+    # "batch" is the output of the training data loader.
+    preds, losses, mae = self.process_batch_supervised(batch)
+    self.log_all(losses, mae, prefix='train_')
+
+    return losses['final'] 
+
+  def validation_step(self, batch, batch_idx):
+    preds, losses, mae = self.process_batch_supervised(batch)
+    self.log_all(losses, mae, prefix='val_')
+  
+  def test_step(self, batch, batch_idx):
+    preds, losses, mae = self.process_batch_supervised(batch)
+    self.log_all(losses, mae, prefix='test_')
+
+       
+    
+# class StripsModel(StripsModelGeneral):
+#   def __init__(self, model_name, *args, **kwargs):
+#     super().__init__( *args, **kwargs)
+#     self.model = timm.create_model(model_name, in_chans=1, num_classes=4)
+#   def forward(self, x):
+#     """get predictions from image batch"""
+#     preds = self.model(x) # preds: logit angle_sin, logit angle_cos, period, logit lumen fraction
+#     preds_sin = 1. - 2*torch.sigmoid(preds[:,0])
+#     preds_cos = 1. - 2*torch.sigmoid(preds[:,1])
+#     preds_direction = 0.5*torch.rad2deg(torch.arctan2(preds_sin, preds_cos))
+#     preds_period = preds[:,2]
+#     preds_lumen_fraction = torch.sigmoid(preds[:,3]) #lumen fraction is between 0 and 1, so we take sigmoid fo this
+#     return preds_direction, preds_period, preds_lumen_fraction
+
+# class StripsModelAngle1(StripsModelGeneral):
+#   def __init__(self, model_name, *args, **kwargs):
+#     super().__init__( *args, **kwargs)
+#     self.model = timm.create_model(model_name, in_chans=1, num_classes=3)
+#   def forward(self, x):
+#     """get predictions from image batch"""
+#     preds = self.model(x) # preds: logit angle_sin, logit angle
+#     preds_direction = torch.pi * torch.sigmoid(preds[:,0])
+#     preds_period = preds[:,1]
+#     preds_lumen_fraction = torch.sigmoid(preds[:,2]) #lumen fraction is between 0 and 1, so we take sigmoid fo this
+#     return preds_direction, preds_period, preds_lumen_fraction       
+        

angle_calculation/classic.py

@@ -0,0 +1,349 @@
+import numpy as np
+from scipy import signal
+from scipy import ndimage
+from scipy.fftpack import next_fast_len
+from skimage.transform import rotate
+from skimage._shared.utils import convert_to_float
+from skimage.transform import warp
+import matplotlib.pyplot as plt
+import cv2
+from copy import deepcopy
+
+def get_directional_std(image, theta=None,*, preserve_range=False):
+    
+    if image.ndim != 2:
+        raise ValueError('The input image must be 2-D')
+    if theta is None:
+        theta = np.arange(180)
+
+    image = convert_to_float(image.copy(), preserve_range) #TODO: needed?
+
+    shape_min = min(image.shape)
+    img_shape = np.array(image.shape)
+
+    # Crop image to make it square
+    slices = tuple(slice(int(np.ceil(excess / 2)),
+                      int(np.ceil(excess / 2) + shape_min))
+                if excess > 0 else slice(None)
+                for excess in (img_shape - shape_min))
+    image = image[slices]
+    shape_min = min(image.shape)
+    img_shape = np.array(image.shape)
+
+    radius = shape_min // 2
+    coords = np.array(np.ogrid[:image.shape[0], :image.shape[1]],
+                      dtype=object)
+    dist = ((coords - img_shape // 2) ** 2).sum(0)
+    outside_reconstruction_circle = dist > radius ** 2
+    image[outside_reconstruction_circle] = 0
+
+    valid_square_slice = slice(int(np.ceil(radius*(1-1/np.sqrt(2)))), int(np.ceil(radius*(1+1/np.sqrt(2)))) )
+
+    # padded_image is always square
+    if image.shape[0] != image.shape[1]:
+        raise ValueError('padded_image must be a square')
+    center = image.shape[0] // 2
+    result = np.zeros(len(theta))
+
+    for i, angle in enumerate(np.deg2rad(theta)):
+        cos_a, sin_a = np.cos(angle), np.sin(angle)
+        R = np.array([[cos_a, sin_a, -center * (cos_a + sin_a - 1)],
+                      [-sin_a, cos_a, -center * (cos_a - sin_a - 1)],
+                      [0, 0, 1]])
+        rotated = warp(image, R, clip=False)
+        result[i] = rotated[valid_square_slice, valid_square_slice].std(axis=0).mean()
+    return result
+
+def acf2d(x, nlags=None):
+  xo = x - x.mean(axis=0)
+  n = len(x)
+  if nlags is None:
+    nlags = n -1
+  lag_len = nlags
+
+  xi = np.arange(1, n + 1)
+  d = np.expand_dims(np.hstack((xi, xi[:-1][::-1])),1)
+
+  nobs = len(xo)
+  n = next_fast_len(2 * nobs + 1)
+  Frf = np.fft.fft(xo, n=n, axis=0)
+
+  acov = np.fft.ifft(Frf * np.conjugate(Frf), axis=0)[:nobs] / d[nobs - 1 :]
+  acov = acov.real
+  ac = acov[: nlags + 1] / acov[:1]
+  return ac
+
+def get_period(acf_table, n_samples=50):
+    #TODO: use peak heights to select best candidates. use std to eliminate outliers
+    period_candidates = []
+    period_candidates_hights = []
+    for i in np.random.randint(0, acf_table.shape[1], min(acf_table.shape[1], n_samples)):
+        peaks = signal.find_peaks(acf_table[:,i])[0]
+        if len(peaks) == 0:
+            continue
+        peak_idx = peaks[0]
+        period_candidates.append(peak_idx)
+        period_candidates_hights.append(acf_table[peak_idx,i])
+    period_candidates = np.array(period_candidates)
+    period_candidates_hights = np.array(period_candidates_hights)
+    
+    if len(period_candidates) == 0:
+        return np.nan, np.nan
+    elif len(period_candidates) == 1:
+        return period_candidates[0], np.nan
+    q1, q3 = np.quantile(period_candidates, [0.25, 0.75])
+    candidates_std = np.std(period_candidates[(period_candidates>=q1)&(period_candidates<=q3)])
+    # return period_candidates, period_candidates_hights
+    return np.median(period_candidates), candidates_std
+
+def get_rotation_with_confidence(padded_image, blur_size=4, make_plots=True):
+  std_by_angle = get_directional_std(cv2.blur(padded_image, (blur_size,blur_size)))
+  rotation_angle = np.argmin(std_by_angle)
+
+  rotation_quality = 1 - np.min(std_by_angle)/np.median(std_by_angle)
+  if make_plots:
+    plt.plot(std_by_angle)
+    plt.axvline(rotation_angle, c='k')
+    plt.title(f'quality: {rotation_quality:0.2f}')
+  return rotation_angle, rotation_quality
+
+def calculate_autocorrelation(oriented_img, blur_kernel=(7,1), make_plots=True):
+  autocorrelation = acf2d(cv2.blur(oriented_img.T, blur_kernel))
+  if make_plots:
+    fig, axs = plt.subplots(ncols=2, figsize=(12,6))
+    axs[0].imshow(autocorrelation)
+    axs[1].plot(autocorrelation.sum(axis=1))
+  return autocorrelation
+
+def get_period_with_confidence(autocorrelation_tab, n_samples=30):
+    period, period_std = get_period(autocorrelation_tab, n_samples=n_samples)
+    if period_std == np.nan:
+        period_confidence = 0.001
+    else:
+        period_confidence = period/(period+2*period_std)
+    return period, period_confidence
+
+def calculate_white_fraction(img, blur_size=4, make_plots=True): #TODO: add mask
+  blurred = cv2.blur(img, (blur_size, blur_size))
+  blurred_sum = blurred.sum(axis=0)
+  lower, upper = np.quantile(blurred_sum, [0.15, 0.85])
+  sign = blurred_sum > (lower+upper)/2
+
+  sign_change = sign[:-1] != sign[1:]
+  sign_change_indices = np.where(sign_change)[0]
+
+  if len(sign_change_indices) >= 2 + (sign[-1] == sign[0]):
+    cut_first = sign_change_indices[0]+1
+
+    if sign[-1] == sign[0]:
+      cut_last = sign_change_indices[-2]
+    else:
+      cut_last = sign_change_indices[-1]
+
+    white_fraction = np.mean(sign[cut_first:cut_last])
+  else:
+    white_fraction = np.nan
+    cut_first, cut_last = None, None 
+  if make_plots:
+      fig, axs = plt.subplots(ncols=3, figsize=(16,6))
+      blurred_sum_normalized = blurred_sum - blurred_sum.min()
+      blurred_sum_normalized /= blurred_sum_normalized.max()
+      axs[0].plot(blurred_sum_normalized)
+      axs[0].plot(sign)
+      axs[1].plot(blurred_sum_normalized[cut_first:cut_last])
+      axs[1].plot(sign[cut_first:cut_last])
+      axs[2].imshow(img, cmap='gray')
+      for i, idx in enumerate(sign_change_indices):
+        plt.axvline(idx, c=['r', 'lime'][i%2])
+      fig.suptitle(f'fraction: {white_fraction:0.2f}')
+
+  return white_fraction
+
+def process_img_crop(img, nm_per_px=1, make_plots=False, return_extra=False):
+
+  # image must be square
+  assert img.shape[0] == img.shape[1]
+  crop_size = img.shape[0]
+
+  # find orientation
+  rotation_angle, rotation_quality = get_rotation_with_confidence(img, blur_size=4, make_plots=make_plots)
+
+  # rotate and crop image
+  crop_margin = int((1 - 1/np.sqrt(2))*crop_size*0.5)
+  oriented_img = rotate(img, -rotation_angle)[2*crop_margin:-crop_margin, crop_margin:-crop_margin]
+
+  # calculate autocorrelation
+  autocorrelation = calculate_autocorrelation(oriented_img, blur_kernel=(7,1), make_plots=make_plots)
+
+  # find period
+  period, period_confidence = get_period_with_confidence(autocorrelation)
+  if make_plots:
+    print(f'period: {period}, confidence: {period_confidence}')
+
+  # find white fraction
+  white_fraction = calculate_white_fraction(oriented_img, make_plots=make_plots)
+  white_width = white_fraction*period
+  
+  result = {
+      'direction': rotation_angle,
+      'direction confidence': rotation_quality,
+      'period': period*nm_per_px,
+      'period confidence': period_confidence,
+      'lumen width': white_width*nm_per_px
+  }
+  if return_extra:
+    result['extra'] = {
+          'autocorrelation': autocorrelation,
+          'oriented_img': oriented_img
+      }
+
+  return result
+
+def get_top_k(a, k):
+  ind = np.argpartition(a, -k)[-k:]
+  return a[ind]
+
+def get_crops(img, distance_map, crop_size, N_sample):
+  crop_r= np.sqrt(2)*crop_size / 2
+  possible_positions_y, possible_positions_x  = np.where(distance_map >= crop_r)
+  no_edge_mask = (possible_positions_y>crop_r) & \
+   (possible_positions_x>crop_r) & \
+   (possible_positions_y<(distance_map.shape[0]-crop_r)) & \
+   (possible_positions_x<(distance_map.shape[1]-crop_r))
+
+  possible_positions_x = possible_positions_x[no_edge_mask]
+  possible_positions_y = possible_positions_y[no_edge_mask]
+  N_available = len(possible_positions_x)
+  positions_indices = np.random.choice(np.arange(N_available), min(N_sample, N_available), replace=False)
+
+  for idx in positions_indices:
+    yield img[possible_positions_y[idx]-crop_size//2:possible_positions_y[idx]+crop_size//2,possible_positions_x[idx]-crop_size//2:possible_positions_x[idx]+crop_size//2].copy()
+
+def sliced_mean(x, slice_size):
+    cs_y = np.cumsum(x, axis=0)
+    cs_y = np.concatenate((np.zeros((1, cs_y.shape[1]), dtype=cs_y.dtype), cs_y), axis=0)
+    slices_y = (cs_y[slice_size:] - cs_y[:-slice_size])/slice_size
+    cs_xy = np.cumsum(slices_y, axis=1)
+    cs_xy = np.concatenate((np.zeros((cs_xy.shape[0], 1), dtype=cs_xy.dtype), cs_xy), axis=1)
+    slices_xy = (cs_xy[:,slice_size:] - cs_xy[:,:-slice_size])/slice_size
+    return slices_xy
+
+def sliced_var(x, slice_size):
+    x = x.astype('float64')
+    return sliced_mean(x**2, slice_size) - sliced_mean(x, slice_size)**2
+    
+def select_samples(granum_image, granum_mask, crop_size=96, n_samples=64, granum_fraction_min=1.0, variance_p=2):
+    granum_occupancy = sliced_mean(granum_mask, crop_size)
+    possible_indices = np.stack(np.where(granum_occupancy >= granum_fraction_min), axis=1)
+    
+    if variance_p == 0:
+        p = np.ones(len(possible_indices))
+    else:
+        variance_map = sliced_var(granum_image, crop_size)
+        p = variance_map[possible_indices[:,0], possible_indices[:,1]]**variance_p
+    p /= np.sum(p)
+    
+    chosen_indices = np.random.choice(
+        np.arange(len(possible_indices)),
+        min(len(possible_indices), n_samples),
+        replace=False,
+        p = p
+    )
+
+    crops = []
+    for crop_idx, idx in enumerate(chosen_indices):
+        crops.append(
+            granum_image[
+                possible_indices[idx,0]:possible_indices[idx,0]+crop_size,
+                possible_indices[idx,1]:possible_indices[idx,1]+crop_size
+            ]
+        )
+    return np.array(crops)
+
+def calculate_distance_map(mask):
+    padded = np.pad(mask, pad_width=1, mode='constant', constant_values=False)
+    distance_map_padded = ndimage.distance_transform_edt(padded)
+    return distance_map_padded[1:-1,1:-1]
+
+
+def measure_object(
+  img, mask,
+  nm_per_px=1, n_tries = 3,
+  direction_thr_min = 0.07, direction_thr_enough = 0.1,
+  crop_size = 200,
+  **kwargs): 
+
+    distance_map = calculate_distance_map(mask)
+    crop_size = min(crop_size, int(min(get_top_k(distance_map.flatten(), n_tries)*0.5**0.5)))
+  
+    direction_confidence = 0
+    best_stripes_data = {}
+    for i, img_crop in enumerate(get_crops(img, distance_map, crop_size, n_tries)):
+        stripes_data = process_img_crop(img_crop, nm_per_px=nm_per_px)
+        if stripes_data['direction confidence'] >= direction_confidence:
+            best_stripes_data = deepcopy(stripes_data)
+            direction_confidence = stripes_data['direction confidence']
+        if direction_confidence > direction_thr_enough:
+            break
+
+    result = best_stripes_data
+
+    if direction_confidence >= direction_thr_min:
+    
+        mask_oriented = rotate(mask, 90-result['direction'], resize=True).astype('bool')
+        idx_begin_x, idx_end_x = np.where(np.any(mask_oriented, axis=0))[0][np.array([0, -1])]
+        idx_begin_y, idx_end_y = np.where(np.any(mask_oriented, axis=1))[0][np.array([0, -1])]
+        result['mask_oriented'] = mask_oriented[idx_begin_y:idx_end_y, idx_begin_x:idx_end_x]
+        result['img_oriented'] = rotate(img, 90-result['direction'], resize=True)[idx_begin_y:idx_end_y, idx_begin_x:idx_end_x]
+
+  #   measurements = measure_granum_shape(result['mask_oriented'], nm_per_px=nm_per_px, oriented=True)
+  # else:
+  #   measurements = measure_granum_shape(mask, nm_per_px=nm_per_px, oriented=False)
+  
+  # result.update(**measurements)
+  # N_layers = result['height'] / result['period']
+  # if np.isfinite(N_layers):
+  #   N_layers = round(N_layers)
+
+    return result
+
+# def measure_object(
+#     img, mask,
+#     nm_per_px=1, n_tries = 3,
+#     direction_thr_min = 0.07, direction_thr_enough = 0.1,
+#     crop_size = 200,
+#     **kwargs): 
+
+#     distance_map = calculate_distance_map(mask)
+#     crop_size = min(crop_size, int((min(get_top_k(distance_map.flatten(), n_tries)*0.5)**0.5)))
+  
+#     direction_confidence = 0
+#     best_stripes_data = {}
+#     for i, img_crop in enumerate(select_samples(img, mask, crop_size=crop_size, n_samples=n_tries)):
+#         stripes_data = process_img_crop(img_crop, nm_per_px=nm_per_px)
+#         if stripes_data['direction_confidence'] >= direction_confidence:
+#             best_stripes_data = deepcopy(stripes_data)
+#             direction_confidence = stripes_data['direction_confidence']
+#         if direction_confidence > direction_thr_enough:
+#             break
+
+#     result = best_stripes_data
+
+#     if direction_confidence >= direction_thr_min:
+    
+#         mask_oriented = rotate(mask, 90-result['direction'], resize=True).astype('bool')
+#         idx_begin_x, idx_end_x = np.where(np.any(mask_oriented, axis=0))[0][np.array([0, -1])]
+#         idx_begin_y, idx_end_y = np.where(np.any(mask_oriented, axis=1))[0][np.array([0, -1])]
+#         result['mask_oriented'] = mask_oriented[idx_begin_y:idx_end_y, idx_begin_x:idx_end_x]
+#         result['img_oriented'] = rotate(img, 90-result['direction'], resize=True)[idx_begin_y:idx_end_y, idx_begin_x:idx_end_x]
+
+#   #   measurements = measure_granum_shape(result['mask_oriented'], nm_per_px=nm_per_px, oriented=True)
+#   # else:
+#   #   measurements = measure_granum_shape(mask, nm_per_px=nm_per_px, oriented=False)
+  
+#   # result.update(**measurements)
+#   # N_layers = result['height'] / result['period']
+#   # if np.isfinite(N_layers):
+#   #   N_layers = round(N_layers)
+
+#     return result #{**measurements, **best_stripes_data, 'N layers': N_layers}

angle_calculation/envelope_correction.py

@@ -0,0 +1,34 @@
+import numpy as np
+import scipy
+
+def detect_boundaries(mask, axis):
+    # calculate the boundaries of the mask
+    #axis = 0 results in x_from, x_to
+    #axis = 1 results in y_from, y_to
+
+
+    sum = mask.sum(axis=axis)
+
+    ind_from = min(sum.nonzero()[0])
+    ind_to = max(sum.nonzero()[0])
+    return ind_from, ind_to
+
+def area(mask):
+    x1, y1 = detect_boundaries(mask, 0)
+    a = y1 - x1
+    x2, y2 = detect_boundaries(mask, 1)
+    b = y2 - x2
+
+    return (a * b, x1, y1, x2, y2)
+
+def calculate_best_angle_from_mask(mask, angles=np.arange(-10,10,0.5)):
+        areas = []
+        for angle in angles:  
+            rotated_mask = scipy.ndimage.rotate(mask, angle, reshape=True, order = 0)  # order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+            this_area = area(rotated_mask)
+            areas.append(this_area[0])
+            
+        best_angle = angles[np.argmin(areas)]
+        return best_angle
+
+        

angle_calculation/granum_utils.py

@@ -0,0 +1,80 @@
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+from scipy import ndimage
+
+from dataclasses import dataclass
+from typing import Any, List
+from zipfile import ZipFile
+
+def add_text(image: Image.Image, text: str, location=(0.5, 0.5), color='red', size=40) -> Image.Image:
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default(size=size)
+    draw.text((int(image.size[0]*location[0]), int(image.size[1]*location[1])), text, font=font, fill=color)
+    return image
+
+
+def select_unique_mask(mask):
+    """if mask consists of multiple parts, select the largest"""
+    blobs = ndimage.label(mask)[0]
+    blob_labels, blob_sizes = np.unique(blobs, return_counts=True)
+    best_blob_label = blob_labels[1:][np.argmax(blob_sizes[1:])]
+    return blobs == best_blob_label
+
+def object_slice(mask, margin=128):
+    rows = np.any(mask, axis=1)
+    cols = np.any(mask, axis=0)
+    row_min, row_max = np.where(rows)[0][[0, -1]]
+    col_min, col_max = np.where(cols)[0][[0, -1]]
+
+    # Create a slice object for the bounding box
+    bounding_box_slice = (
+        slice(max(0,row_min-margin), min(row_max + 1+margin, len(rows)+1)),
+        slice(max(0,col_min-margin), min(col_max + 1+margin, len(cols)+1))
+    )
+    
+    return bounding_box_slice
+
+def resize_to(image: Image.Image, s=4032) -> Image.Image:
+    w, h = image.size
+    longest_size = max(h, w)
+        
+    resize_factor = longest_size / s
+        
+    resized_image = image.resize((int(w/resize_factor), int(h/resize_factor)))
+    return resized_image
+
+def rolling_mean(x, window):
+    cs = np.r_[0, np.cumsum(x)]
+    rolling_sum = cs[window:] - cs[:-window]
+    return rolling_sum/window
+
+@dataclass
+class Granum:
+    image: Any = None#Optional[np.ndarray] = None
+    mask: Any = None #Optional[np.ndarray] = None
+    scaler: Any = None
+    nm_per_px: float = float('nan')
+    detection_confidence: float = float('nan')
+
+def zip_files(files: List[str], output_name: str) -> None:
+    with ZipFile(output_name, "w") as zipObj:
+        for file in files:
+            zipObj.write(file)
+
+def filter_boundary_detections(masks, scaler=None):
+    last_index_right = -1 if scaler is None else masks.shape[1]-1-scaler.pad_right
+    last_index_bottom = -1 if scaler is None else masks.shape[2]-1-scaler.pad_bottom
+    doesnt_touch_boundary_mask = ~(np.any(masks[:,0,:] != 0, axis=1) | np.any(masks[:,last_index_right:,:] != 0, axis=(1,2)) | np.any(masks[:,:,0] != 0, axis=1) | np.any(masks[:,:,last_index_bottom:] != 0, axis=(1,2)))
+    return doesnt_touch_boundary_mask
+
+def get_circle_mask(shape, r=None):
+    if isinstance(shape, int):
+        shape = (shape, shape)
+    if r is None:
+        r = min(shape)/2
+    X, Y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+    center_x = shape[1] / 2 - 0.5
+    center_y = shape[0] / 2 - 0.5
+
+    mask = ((X-center_x)**2 + (Y-center_y)**2) >= r**2
+    return mask

angle_calculation/image_transforms.py

@@ -0,0 +1,34 @@
+import numpy as np
+import cv2
+
+def batched_radon(image_batch):
+    batch_size, img_size = image_batch.shape[:2]
+    if batch_size > 512: # limit batch size to 512 because cv2.warpAffine fails for batch> 512
+        return np.concatenate([batched_radon(image_batch[i:i+512]) for i in range(0,batch_size,512)], axis=0)
+    theta = np.arange(180)
+    radon_image = np.zeros((image_batch.shape[0], img_size, len(theta)),
+                           dtype='float32')
+
+    for i, angle in enumerate(theta):
+        M = cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1)
+        rotated = cv2.warpAffine(np.transpose(image_batch, (1, 2, 0)),M,(img_size,img_size))
+        if batch_size == 1: # cv2.warpAffine cancels batch dimension if equal to 1
+          rotated = rotated[:,:, np.newaxis]
+        rotated = np.transpose(rotated, (2, 0, 1))
+        rotated = rotated / np.array(255, dtype='float32')
+        radon_image[:, :, i] = rotated.sum(axis=1)
+    return radon_image
+
+def get_center_crop_coords(height: int, width: int, crop_height: int, crop_width: int):
+    """from https://github.com/albumentations-team/albumentations/blob/master/albumentations/augmentations/crops/functional.py"""
+    y1 = (height - crop_height) // 2
+    y2 = y1 + crop_height
+    x1 = (width - crop_width) // 2
+    x2 = x1 + crop_width
+    return x1, y1, x2, y2
+
+def center_crop(img: np.ndarray, crop_height: int, crop_width: int):
+    height, width = img.shape[:2]
+    x1, y1, x2, y2 = get_center_crop_coords(height, width, crop_height, crop_width)
+    img = img[y1:y2, x1:x2]
+    return img

angle_calculation/sampling.py

@@ -0,0 +1,142 @@
+from PIL import Image
+from typing import Tuple
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy import ndimage
+import torch
+from torchvision.transforms import functional as tvf
+
+from pathlib import Path
+
+def sliced_mean(x, slice_size):
+    cs_y = np.cumsum(x, axis=0)
+    cs_y = np.concatenate((np.zeros((1, cs_y.shape[1]), dtype=cs_y.dtype), cs_y), axis=0)
+    slices_y = (cs_y[slice_size:] - cs_y[:-slice_size])/slice_size
+    cs_xy = np.cumsum(slices_y, axis=1)
+    cs_xy = np.concatenate((np.zeros((cs_xy.shape[0], 1), dtype=cs_xy.dtype), cs_xy), axis=1)
+    slices_xy = (cs_xy[:,slice_size:] - cs_xy[:,:-slice_size])/slice_size
+    return slices_xy
+
+def sliced_var(x, slice_size):
+    x = x.astype('float64')
+    return sliced_mean(x**2, slice_size) - sliced_mean(x, slice_size)**2
+
+def calculate_local_variance(img, var_window):
+    """return local variance map with the same size as input image"""
+    var = sliced_var(img, var_window)
+
+    left_pad = var_window // 2 -1
+    right_pad = var_window -1 - left_pad
+    var_padded = np.pad(
+        var,
+        pad_width=(
+            (left_pad,right_pad),
+            (left_pad,right_pad)
+        ))
+    return var_padded
+
+def get_crop_batch(img: np.ndarray, mask: np.ndarray, crop_size=96, crop_scales=np.geomspace(0.5, 2, 7), samples_per_scale=32, use_variance_threshold=False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Generate a batch of cropped images from an input image and corresponding mask, at various scales and rotations.
+
+    Parameters
+    ----------
+    img : np.ndarray
+        The input image from which crops are generated.
+    mask : np.ndarray
+        The binary mask indicating the region of interest in the image.
+    crop_size : int, optional
+        The size of the square crop.
+    crop_scales : np.ndarray, optional
+        An array of scale factors to apply to the crop size.
+    samples_per_scale : int, optional
+        Number of samples to generate per scale factor.
+    use_variance_threshold : bool, optional
+        Flag to use variance thresholding for selecting crop locations.
+
+    Returns
+    -------
+    Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        A tuple containing the tensor of crops, their rotation angles, and scale factors.
+    """
+    
+    # pad
+    pad_size = int(np.ceil(0.5*crop_size*max(crop_scales)*(np.sqrt(2)-1)))
+    img_padded = np.pad(img, pad_size)
+    mask_padded = np.pad(mask, pad_size)
+
+    # distance map
+    distance_map_padded = ndimage.distance_transform_edt(mask_padded)
+    # TODO: adjust scales and samples_per_scale
+    
+    if use_variance_threshold:
+        variance_window = min(crop_size//2, min(img.shape))
+        variance_map_padded = np.pad(calculate_local_variance(img, variance_window), pad_size)
+        variance_median = np.ma.median(np.ma.masked_where(distance_map_padded<0.5*variance_window, variance_map_padded))
+        variance_mask = variance_map_padded >= variance_median
+    else:
+        variance_mask = np.ones_like(mask_padded)
+    
+    # initilize output
+    crops_granum = []
+    angles_granum = []
+    scales_granum = []
+    # loop over scales
+    for scale in crop_scales: 
+        half_crop_size_scaled = int(np.floor(scale*0.5*crop_size)) # half of crop size after scaling
+        crop_pad = int(np.ceil((np.sqrt(2) - 1)*half_crop_size_scaled)) # pad added in order to allow rotation
+        half_crop_size_external = half_crop_size_scaled + crop_pad # size of "external crop" which will be rotated
+
+        possible_indices = np.stack(np.where(variance_mask & (distance_map_padded >= 2*half_crop_size_scaled)), axis=1)
+        if len(possible_indices) == 0:
+            continue
+        chosen_indices = np.random.choice(np.arange(len(possible_indices)), min(len(possible_indices), samples_per_scale), replace=False)
+
+        crops = [
+            img_padded[y-half_crop_size_external:y+half_crop_size_external, x-half_crop_size_external:x+half_crop_size_external] for y, x in possible_indices[chosen_indices]
+        ]
+
+        # rotate
+        rotation_angles = np.random.rand(len(crops))*180 - 90
+        crops = [
+            ndimage.rotate(crop, angle, reshape=False)[crop_pad:-crop_pad,crop_pad:-crop_pad] for crop, angle in zip(crops, rotation_angles)
+        ]
+        # add to output
+        crops_granum.append(tvf.resize(torch.tensor(np.array(crops)), (crop_size,crop_size),antialias=True)) # resize crops to crop_size
+        angles_granum.extend(rotation_angles.tolist())
+        scales_granum.extend([scale]*len(crops))
+
+    if len(angles_granum) == 0:
+        return [], [], []
+    
+    crops_granum = torch.concat(crops_granum)
+    angles_granum = torch.tensor(angles_granum, dtype=torch.float)
+    scales_granum = torch.tensor(scales_granum, dtype=torch.float)
+    
+    return crops_granum, angles_granum, scales_granum
+
+def get_crop_batch_from_path(img_path, mask_path=None, use_variance_threshold=False):
+    """
+    Load an image and its mask from file paths and generate a batch of cropped images.
+
+    Parameters
+    ----------
+    img_path : str
+        Path to the input image.
+    mask_path : str, optional
+        Path to the binary mask image. If None, assumes mask path by replacing image extension with '.npy'.
+    use_variance_threshold : bool, optional
+        Flag to use variance thresholding for selecting crop locations.
+
+    Returns
+    -------
+    Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        A tuple containing the tensor of crops, their rotation angles, and scale factors, obtained from the specified image path.
+    """
+    if mask_path is None:
+        mask_path = str(Path(img_path).with_suffix('.npy'))
+    mask = np.load(mask_path)
+    img = np.array(Image.open(img_path))[:,:,0]
+    
+    return get_crop_batch(img, mask, use_variance_threshold=use_variance_threshold)
+    

app.py

@@ -0,0 +1,602 @@
+import os
+import shutil
+import time
+import uuid
+from datetime import datetime
+from decimal import Decimal
+
+import gradio as gr
+import matplotlib.pyplot as plt
+
+from settings import DEMO
+
+plt.switch_backend("agg")  # fix for "RuntimeError: main thread is not in main loop"
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+from model import GranaAnalyser
+
+ga = GranaAnalyser(
+    "weights/yolo/20240604_yolov8_segm_ABRCR1_all_train4_best.pt",
+    "weights/AS_square_v16.ckpt",
+    "weights/period_measurer_weights-1.298_real_full-fa12970.ckpt",
+)
+
+
+def calc_ratio(pixels, nano):
+    """
+    Calculates ratio of pixels to nanometers and returns as str to populate ratio_input
+    :param pixels:
+    :param nano:
+    :return:
+    """
+    if not (pixels and nano):
+        pass
+    else:
+        res = pixels / nano
+        return res
+
+
+# https://jakevdp.github.io/PythonDataScienceHandbook/05.13-kernel-density-estimation.html
+def KDE(dataset, h):
+    # the Kernel function
+    def K(x):
+        return np.exp(-(x ** 2) / 2) / np.sqrt(2 * np.pi)
+
+    n_samples = dataset.size
+
+    x_range = dataset  # x-value range for plotting KDEs
+
+    total_sum = 0
+    # iterate over datapoints
+    for i, xi in enumerate(dataset):
+        total_sum += K((x_range - xi) / h)
+
+        y_range = total_sum / (h * n_samples)
+
+    return y_range
+
+
+def prepare_files_for_download(
+    dir_name,
+    grana_data,
+    aggregated_data,
+    detection_visualizations_dict,
+    images_grana_dict,
+):
+    """
+    Save and zip files for download
+    :param dir_name:
+    :param grana_data: DataFrame containing all grana measurements
+    :param aggregated_data: dict containing aggregated measurements
+    :return:
+    """
+    dir_to_zip = f"{dir_name}/to_zip"
+
+    # raw data
+    grana_data_csv_path = f"{dir_to_zip}/grana_raw_data.csv"
+    grana_data.to_csv(grana_data_csv_path, index=False)
+
+    # aggregated measurements
+    aggregated_csv_path = f"{dir_to_zip}/grana_aggregated_data.csv"
+    aggregated_data.to_csv(aggregated_csv_path)
+
+    # annotated pictures
+    masked_images_dir = f"{dir_to_zip}/annotated_images"
+    os.makedirs(masked_images_dir)
+    for img_name, img in detection_visualizations_dict.items():
+        filename_split = img_name.split(".")
+        extension = filename_split[-1]
+        filename = ".".join(filename_split[:-1])
+        filename = f"{filename}_annotated.{extension}"
+        img.save(f"{masked_images_dir}/{filename}")
+
+    # single_grana images
+    grana_images_dir = f"{dir_to_zip}/single_grana_images"
+    os.makedirs(grana_images_dir)
+    org_images_dict = pd.Series(
+        grana_data["source image"].values, index=grana_data["granum ID"]
+    ).to_dict()
+    for img_name, img in images_grana_dict.items():
+        org_filename = org_images_dict[img_name]
+        org_filename_split = org_filename.split(".")
+        org_filename_no_ext = ".".join(org_filename_split[:-1])
+        img_name_ext = f"{org_filename_no_ext}_granum_{str(img_name)}.png"
+        img.save(f"{grana_images_dir}/{img_name_ext}")
+
+    # zip all files
+    date_str = datetime.today().strftime("%Y-%m-%d")
+    zip_name = f"GRANA_results_{date_str}"
+    zip_path = f"{dir_name}/{zip_name}"
+    shutil.make_archive(zip_path, "zip", dir_to_zip)
+
+    # delete to_zip dir
+    zip_dir_path = os.path.join(os.getcwd(), dir_to_zip)
+    shutil.rmtree(zip_dir_path)
+
+    download_file_path = f"{zip_path}.zip"
+    return download_file_path
+
+
+def show_info_on_submit(s):
+    return (
+        gr.Button(interactive=False),
+        gr.Button(interactive=False),
+        gr.Row(visible=True),
+        gr.Row(visible=False),
+    )
+
+
+def load_css():
+    with open("styles.css", "r") as f:
+        css_content = f.read()
+    return css_content
+
+
+primary_hue = gr.themes.Color(
+    c50="#e1f8ee",
+    c100="#b7efd5",
+    c200="#8de6bd",
+    c300="#63dda5",
+    c400="#39d48d",
+    c500="#27b373",
+    c600="#1e8958",
+    c700="#155f3d",
+    c800="#0c3522",
+    c900="#030b07",
+    c950="#000",
+)
+
+
+theme = gr.themes.Default(
+    primary_hue=primary_hue,
+    font=[gr.themes.GoogleFont("Ubuntu"), "ui-sans-serif", "system-ui", "sans-serif"],
+)
+
+
+def draw_violin_plot(y, ylabel, title):
+    # only generate plot for 3 or more values
+    if y.count() < 3:
+        return None
+
+    # Colors
+    RED_DARK = "#850e00"
+    DARK_GREEN = "#0c3522"
+    BRIGHT_GREEN = "#8de6bd"
+
+    # Create jittered version of "x" (which is only 1)
+    x_jittered = []
+    kde = KDE(y, (y.max() - y.min()) / y.size / 2)
+    kde = kde / kde.max() * 0.2
+    for y_val in kde:
+        x_jittered.append(1 + np.random.uniform(-y_val, y_val, 1))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+    ax.scatter(x=x_jittered, y=y, s=20, alpha=0.4, c=DARK_GREEN)
+
+    violins = ax.violinplot(
+        y,
+        widths=0.45,
+        bw_method="silverman",
+        showmeans=False,
+        showmedians=False,
+        showextrema=False,
+    )
+
+    # change violin color
+    for pc in violins["bodies"]:
+        pc.set_facecolor(BRIGHT_GREEN)
+
+    # add a boxplot to ax
+    # but make the whiskers length equal to 1 SD, i.e. in the proportion of the IQ range, but this length should start from the mean but be visible from the box boundary
+    lower = np.mean(y) - 1 * np.std(y)
+    upper = np.mean(y) + 1 * np.std(y)
+
+    medianprops = dict(linewidth=1, color="black", solid_capstyle="butt")
+    boxplot_stats = [
+        {
+            "med": np.median(y),
+            "q1": np.percentile(y, 25),
+            "q3": np.percentile(y, 75),
+            "whislo": lower,
+            "whishi": upper,
+        }
+    ]
+
+    ax.bxp(
+        boxplot_stats,  # data for the boxplot
+        showfliers=False,  # do not show the outliers beyond the caps.
+        showcaps=True,  # show the caps
+        medianprops=medianprops,
+    )
+
+    # Add mean value point
+    ax.scatter(1, y.mean(), s=30, color=RED_DARK, zorder=3)
+
+    ax.set_xticks([])
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+    fig.tight_layout()
+
+    return fig
+
+
+def transform_aggregated_results_table(results_dict):
+    MEASUREMENT_HEADER = "measurement [unit]"
+    VALUE_HEADER = "value +-SD"
+
+    def get_value_str(value, std):
+        if np.isnan(value) or np.isnan(std):
+            return "-"
+        value_str = str(Decimal(str(value)).quantize(Decimal("0.01")))
+        std_str = str(Decimal(str(std)).quantize(Decimal("0.01")))
+        return f"{value_str} +-{std_str}"
+
+    def append_to_dict(new_key, old_val_key, old_sd_key):
+        aggregated_dict[MEASUREMENT_HEADER].append(new_key)
+        value_str = get_value_str(results_dict[old_val_key], results_dict[old_sd_key])
+        aggregated_dict[VALUE_HEADER].append(value_str)
+
+    aggregated_dict = {MEASUREMENT_HEADER: [], VALUE_HEADER: []}
+
+    # area
+    append_to_dict("area [nm^2]", "area nm^2", "area nm^2 std")
+
+    # perimeter
+    append_to_dict("perimeter [nm]", "perimeter nm", "perimeter nm std")
+
+    # diameter
+    append_to_dict("diameter [nm]", "diameter nm", "diameter nm std")
+
+    # height
+    append_to_dict("height [nm]", "height nm", "height nm std")
+
+    # number of layers
+    append_to_dict("number of thylakoids", "Number of layers", "Number of layers std")
+
+    # SRD
+    append_to_dict("SRD [nm]", "period nm", "period nm std")
+
+    # GSI
+    append_to_dict("GSI", "GSI", "GSI std")
+
+    # N grana
+    aggregated_dict[MEASUREMENT_HEADER].append("number of grana")
+    aggregated_dict[VALUE_HEADER].append(str(int(results_dict["N grana"])))
+
+    return aggregated_dict
+
+
+def rename_columns_in_results_table(results_table):
+    column_names = {
+        "Granum ID": "granum ID",
+        "File name": "source image",
+        "area nm^2": "area [nm^2]",
+        "perimeter nm": "perimeter [nm]",
+        "diameter nm": "diameter [nm]",
+        "height nm": "height [nm]",
+        "Number of layers": "number of thylakoids",
+        "period nm": "SRD [nm]",
+        "period SD nm": "SRD SD [nm]",
+    }
+    results_table = results_table.rename(columns=column_names)
+    return results_table
+
+
+with gr.Blocks(css=load_css(), theme=theme) as demo:
+
+    svg = """
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 30.73 33.38">
+  <defs>
+    <style>
+      .cls-1 {
+        fill: #27b373;
+        stroke-width: 0px;
+      }
+    </style>
+  </defs>
+  <path class="cls-1" d="M19.69,11.73h-3.22c-2.74,0-4.96,2.22-4.96,4.96h0c0,2.74,2.22,4.96,4.96,4.96h3.43c.56,0,1,.51.89,1.09-.08.43-.49.72-.92.72h-8.62c-.74,0-1.34-.6-1.34-1.34v-10.87c0-.74.6-1.34,1.34-1.34h13.44c2.73,0,4.95-2.22,4.95-4.95h0c0-2.75-2.22-4.97-4.96-4.97h-13.85C4.85,0,0,4.85,0,10.83v11.71c0,5.98,4.85,10.83,10.83,10.83h9.07c5.76,0,10.49-4.52,10.81-10.21.35-6.29-4.72-11.44-11.02-11.44ZM19.9,31.4h-9.07c-4.89,0-8.85-3.96-8.85-8.85v-11.71C1.98,5.95,5.95,1.98,10.83,1.98h13.81c1.64,0,2.97,1.33,2.97,2.97h0c0,1.65-1.33,2.97-2.96,2.97h-13.4c-1.83,0-3.32,1.49-3.32,3.32v10.87c0,1.83,1.49,3.32,3.32,3.32h8.56c1.51,0,2.83-1.12,2.97-2.62.16-1.72-1.2-3.16-2.88-3.16h-3.52c-1.64,0-2.97-1.33-2.97-2.97h0c0-1.64,1.33-2.97,2.97-2.97h3.34c4.83,0,8.9,3.81,9.01,8.64s-3.9,9.04-8.84,9.04Z"/>
+  <path class="cls-1" d="M19.9,29.41h-9.07c-3.79,0-6.87-3.07-6.87-6.87v-11.71c0-3.79,3.07-6.87,6.87-6.87h13.81c.55,0,.99.44.99.99h0c0,.55-.44.99-.99.99h-13.81c-2.7,0-4.88,2.19-4.88,4.88v11.71c0,2.7,2.19,4.88,4.88,4.88h8.94c2.64,0,4.91-2.05,5-4.7s-2.12-5.05-4.87-5.05h-3.52c-.55,0-.99-.44-.99-.99h0c0-.55.44-.99.99-.99h3.36c3.74,0,6.9,2.92,7.01,6.66.11,3.87-3.01,7.06-6.85,7.06Z"/>
+</svg>
+"""
+
+    gr.HTML(
+        f'<div class="header"><div id="header-logo">{svg}</div><div id="header-text">GRANA<div></div>'
+    )
+
+    with gr.Row(elem_classes="input-row"):  # input
+        with gr.Column():
+            gr.HTML(
+                "<h1>1. Choose images to upload. All the images need to be of the same scale and experimental variant.</h1>"
+            )
+            img_input = gr.File(file_count="multiple")
+
+            gr.HTML("<h1>2. Set the scale of the images for the measurements.</h1>")
+            with gr.Row():
+                with gr.Column():
+                    gr.HTML("Either provide pixel per nanometer ratio...")
+                    ratio_input = gr.Number(
+                        label="pixel per nm", precision=3, step=0.001
+                    )
+
+                with gr.Column():
+                    gr.HTML("...or length of the scale bar in pixels and nanometers.")
+                    pixels_input = gr.Number(label="Length in pixels")
+                    nano_input = gr.Number(label="Length in nanometers")
+
+                    pixels_input.change(
+                        calc_ratio,
+                        inputs=[pixels_input, nano_input],
+                        outputs=ratio_input,
+                    )
+                    nano_input.change(
+                        calc_ratio,
+                        inputs=[pixels_input, nano_input],
+                        outputs=ratio_input,
+                    )
+
+            with gr.Row():
+                clear_btn = gr.ClearButton(img_input, "Clear")
+                submit_btn = gr.Button("Submit", variant="primary")
+
+    with gr.Row(visible=False) as loading_row:
+        with gr.Column():
+            gr.HTML(
+                "<div class='processed-info'>Images are being processed. This may take a while...</div>"
+            )
+
+    with gr.Row(visible=False) as output_row:
+        with gr.Column():
+            gr.HTML(
+                '<div class="results-header">Results</div>'
+                "<p>Full results are a zip file containing:<p>"
+                "<ul>- grana_raw_data.csv: a table with full grana measurements,</ul>"
+                "<ul>- grana_aggregated_data.csv: a table with aggregated measurements,</ul>"
+                '<ul>- directory "annotated_images" with all submitted images with masks on detected grana,</ul>'
+                '<ul>- directory "single_grana_images" with images of all detected grana.</ul>'
+                "<p>Note that GRANA only stores the result files for 1 hour.</p>",
+                elem_classes="input-row",
+            )
+            with gr.Row(elem_classes="input-row"):
+                download_file_out = gr.DownloadButton(
+                    label="Download results",
+                    variant="primary",
+                    elem_classes="margin-bottom",
+                )
+            with gr.Row():
+                gr.HTML(
+                    '<h2 class="title">Annotated images</h2>'
+                    "Gallery of uploaded images with masks of recognized grana structures. "
+                    "Each granum mask is "
+                    "labeled with its number. Note that only fully visible grana in the image are masked."
+                )
+            with gr.Row(elem_classes="margin-bottom"):
+                gallery_out = gr.Gallery(
+                    columns=4,
+                    rows=2,
+                    object_fit="contain",
+                    label="Detection visualizations",
+                    show_download_button=False,
+                )
+
+            with gr.Row(elem_classes="input-row"):
+                gr.HTML(
+                    '<h2 class="title">Aggregated results for all uploaded images</h2>'
+                )
+            with gr.Row(elem_classes=["input-row", "margin-bottom"]):
+                table_out = gr.Dataframe(label="Aggregated data")
+
+            with gr.Row():
+                gr.HTML(
+                    '<h2 class="title">Violin graphs</h2>'
+                    "These graphs present aggregated results for selected structural parameters. "
+                    "The graph for each parameter is only generated if three or more values are available. "
+                    "Each graph "
+                    "displays individual data points, a box plot indicating the first and third quartiles, whiskers "
+                    "marking the standard deviation (SD), the median value (horizontal line on the box plot), "
+                    "the mean value (red dot), and a density plot where the width represents the frequency."
+                )
+
+            with gr.Row():
+                area_plot_out = gr.Plot(label="Area")
+                perimeter_plot_out = gr.Plot(label="Perimeter")
+                gsi_plot_out = gr.Plot(label="GSI")
+
+            with gr.Row(elem_classes="margin-bottom"):
+                diameter_plot_out = gr.Plot(label="Diameter")
+                height_plot_out = gr.Plot(label="Height")
+                srd_plot_out = gr.Plot(label="SRD")
+
+            with gr.Row():
+                gr.HTML(
+                    '<h2 class="title">Recognized and rotated grana structures</h2>'
+                )
+
+            with gr.Row(elem_classes="margin-bottom"):
+                gallery_single_grana_out = gr.Gallery(
+                    columns=4,
+                    rows=2,
+                    object_fit="contain",
+                    label="Single grana images",
+                    show_download_button=False,
+                )
+
+            with gr.Row():
+                gr.HTML(
+                    '<h2 class="title">Full results</h2>'
+                    "Note that structural parameters other than area and perimeter are only calculated for the grana "
+                    "whose direction and/or SRD could be estimated."
+                )
+            with gr.Row():
+                table_full_out = gr.Dataframe(label="Full measurements data")
+
+    submit_btn.click(
+        show_info_on_submit,
+        inputs=[submit_btn],
+        outputs=[submit_btn, clear_btn, loading_row, output_row],
+    )
+
+    def enable_submit():
+        return (
+            gr.Button(interactive=True),
+            gr.Button(interactive=True),
+            gr.Row(visible=False),
+        )
+
+    def gradio_analize_image(images, scale):
+        """
+        Model accepts following parameters:
+        :param images: list of images to be processed, in either tiff or png format
+        :param scale: float, nm to pixel ratio
+
+        Model returns the following objects:
+        - detection_visualizations: list of images with masks to be displayed as gallery and served to download
+        as zip of images
+        - grana_data: dataframe with measurements for each image to be served to download as a csv file
+        - images_grana: list of images with single grana to be served to download as zip of images
+        - aggregated_data: dataframe with aggregated measurements for all images to be displayed as table and served
+        to download as csv
+        """
+
+        # validate that at least one image has been uploaded
+        if images is None or len(images) == 0:
+            raise gr.Error("Please upload at least one image")
+
+        # on demo instance, we limit the number of images to 5
+        if DEMO:
+            if len(images) > 5:
+                raise gr.Error("In demo version it is possible to analyze up to 5 images.")
+
+        # validate that scale has been provided correctly
+        if scale is None or scale == 0:
+            raise gr.Error("Please provide scale. Use dot as decimal separator")
+
+        # validate that all images are png or tiff
+        for image in images:
+            if not image.name.lower().endswith((".png", ".tif", ".jpg", ".jpeg")):
+                raise gr.Error("Only png, tiff, jpg ang jpeg images are supported")
+
+        # clean up previous results
+        # find all directories in current working directory that start with "results_"
+        # that were created more than 1 hour ago and delete them with all contents
+        for directory_name in os.listdir():
+            if directory_name.startswith("results_"):
+                dir_path = os.path.join(os.getcwd(), directory_name)
+                if os.path.isdir(dir_path):
+                    if time.time() - os.path.getctime(dir_path) > 60 * 60:
+                        shutil.rmtree(dir_path)
+
+        # create a directory for results
+        results_dir_name = "results_{uuid}".format(uuid=uuid.uuid4().hex)
+        os.makedirs(results_dir_name)
+        zip_dir_name = f"{results_dir_name}/to_zip"
+        os.makedirs(zip_dir_name)
+
+        # model takes a dict of images, so we need to convert input to list of PIL.PngImagePlugin.PngImageFile or
+        # PIL.TiffImagePlugin.TiffImageFile objects
+        images_dict = {
+            image.name.split("/")[-1]: Image.open(image.name)
+            for i, image in enumerate(images)
+        }
+
+        # model works here
+        (
+            detection_visualizations_dict,
+            grana_data,
+            images_grana_dict,
+            aggregated_data,
+        ) = ga.predict(images_dict, scale)
+        detection_visualizations = list(detection_visualizations_dict.values())
+        images_grana = list(images_grana_dict.values())
+
+        # rearrange aggregated data to be displayed as table
+        aggregated_dict = transform_aggregated_results_table(aggregated_data)
+        aggregated_df_transposed = pd.DataFrame.from_dict(aggregated_dict)
+
+        # rename columns in full results
+        grana_data = rename_columns_in_results_table(grana_data)
+
+        # save files returned by model to disk so they can be retrieved for downloading
+        download_file_path = prepare_files_for_download(
+            results_dir_name,
+            grana_data,
+            aggregated_df_transposed,
+            detection_visualizations_dict,
+            images_grana_dict,
+        )
+
+        # generate plot
+        area_fig = draw_violin_plot(
+            grana_data["area [nm^2]"].dropna(),
+            "Granum area [nm^2]",
+            "Grana areas from all uploaded images",
+        )
+        perimeter_fig = draw_violin_plot(
+            grana_data["perimeter [nm]"].dropna(),
+            "Granum perimeter [nm]",
+            "Grana perimeters from all uploaded images",
+        )
+        gsi_fig = draw_violin_plot(
+            grana_data["GSI"].dropna(),
+            "GSI",
+            "GSI from all uploaded images",
+        )
+        diameter_fig = draw_violin_plot(
+            grana_data["diameter [nm]"].dropna(),
+            "Granum diameter [nm]",
+            "Grana diameters from all uploaded images",
+        )
+        height_fig = draw_violin_plot(
+            grana_data["height [nm]"].dropna(),
+            "Granum height [nm]",
+            "Grana heights from all uploaded images",
+        )
+        srd_fig = draw_violin_plot(
+            grana_data["SRD [nm]"].dropna(), "SRD [nm]", "SRD from all uploaded images"
+        )
+
+        return [
+            gr.Row(visible=True),
+            gr.Row(visible=True),
+            download_file_path,
+            detection_visualizations,
+            aggregated_df_transposed,
+            area_fig,
+            perimeter_fig,
+            gsi_fig,
+            diameter_fig,
+            height_fig,
+            srd_fig,
+            images_grana,
+            grana_data,
+        ]
+
+    submit_btn.click(
+        fn=gradio_analize_image,
+        inputs=[
+            img_input,
+            ratio_input,
+        ],
+        outputs=[
+            loading_row,
+            output_row,
+            # file_download_checkboxes,
+            download_file_out,
+            gallery_out,
+            table_out,
+            area_plot_out,
+            perimeter_plot_out,
+            gsi_plot_out,
+            diameter_plot_out,
+            height_plot_out,
+            srd_plot_out,
+            gallery_single_grana_out,
+            table_full_out,
+        ],
+    ).then(fn=enable_submit, inputs=[], outputs=[submit_btn, clear_btn, loading_row])
+
+demo.launch(
+    share=False, debug=True, server_name="0.0.0.0", allowed_paths=["images/logo.svg"]
+)

grana_detection/mmwrapper.py

@@ -0,0 +1,42 @@
+import torch
+import numpy as np
+from typing import Union, Optional
+from PIL import Image
+from mmdet.apis import DetInferencer
+from ultralytics.engine.results import Results
+import warnings
+
+class MMDetector(DetInferencer):
+    def __call__(
+            self,
+            inputs,
+            ) -> Results:
+        """Call the inferencer as in DetInferencer but for single image.
+
+        Args:
+            inputs (np.ndarray | str): Inputs for the inferencer.
+
+        Returns:
+            Result: yolo-like result
+        """
+
+        ori_inputs = self._inputs_to_list(inputs)
+
+        data = list(self.preprocess(
+            ori_inputs, batch_size=1))[0][1]
+
+        preds = self.forward(data)[0]
+        
+        yolo_result = Results(
+            orig_img=ori_inputs[0], path="", names=[""],
+            boxes=torch.cat((preds.pred_instances.bboxes, preds.pred_instances.scores.unsqueeze(-1), preds.pred_instances.labels.unsqueeze(-1)), dim=1),
+            masks=preds.pred_instances.masks
+        )
+            
+        return yolo_result
+    
+    def predict(self, source: Image.Image, conf=None):
+        """yolo interface"""
+        if conf is not None:
+            warnings.warn(f"confidence value {conf} ignored")
+        return [self.__call__(np.array(source.convert("RGB")))]

model.py

@@ -0,0 +1,629 @@
+import itertools
+import warnings
+from io import BytesIO
+from copy import deepcopy
+from pathlib import Path
+from typing import List, Tuple, Dict, Optional, Any, Union
+from dataclasses import dataclass, field
+
+from PIL import Image
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from scipy import ndimage
+from skimage import measure
+import torchvision.transforms.functional as tvf
+from ultralytics import YOLO
+import torch
+import cv2
+import gradio
+
+import sys, os
+sys.path.append(os.path.abspath('angle_calculation'))
+# from classic import measure_object
+from sampling import get_crop_batch
+from angle_model import PatchedPredictor, StripsModelLumenWidth
+
+from period_calculation.period_measurer import PeriodMeasurer
+# from grana_detection.mmwrapper import MMDetector # mmdet installation in docker is problematic for now
+
+@dataclass
+class Granum:
+    id: Optional[int] = None
+    image: Any = None
+    mask: Any = None
+    scaler: Any = None
+    nm_per_px: float = float('nan')
+    detection_confidence: float = float('nan')
+    img_oriented: Optional[np.ndarray] = None # oriented fragment of the image
+    mask_oriented: Optional[np.ndarray] = None # oriented fragment of the mask
+    measurements: dict = field(default_factory=dict) # dict with grana measurements
+
+class ScalerPadder:
+    """resize and pad image to specific range.
+    minimal_pad: obligatory padding, e.g. required for detector
+    """
+    def __init__(self, target_size=1024, target_short_edge_min=640, minimal_pad=16, pad_to_multiply=32):
+        self.minimal_pad = minimal_pad
+        self.target_size = target_size - 2*self.minimal_pad # detection pad is necessary padding size
+        self.target_short_edge_min = target_short_edge_min - 2*self.minimal_pad
+
+        self.max_size_nm = 6000 # training images covers ~3100 nm
+        self.min_size_nm = 2400 # training images covers ~3100 nm
+        self.pad_to_multiply = pad_to_multiply
+        
+    def transform(self, image: Image.Image, px_per_nm: float=1.298) -> Image.Image:
+        self.original_size = image.size
+        self.original_px_per_nm = px_per_nm
+        w, h = self.original_size
+        longest_size = max(h, w)
+        img_size_nm = longest_size / px_per_nm
+        if img_size_nm > self.max_size_nm:
+            error_message = f'too large image, image size: {img_size_nm:0.1f}nm, max allowed: {self.max_size_nm}nm'
+            # raise ValueError(error_message)
+            # warnings.warn(warning_message)
+            gradio.Warning(error_message)
+            # add_text(image, warning_message, location=(0.1, 0.1), color='blue', size=int(40*longest_size/self.target_size))
+        
+        self.resize_factor = self.target_size / (max(self.min_size_nm, img_size_nm) * px_per_nm)
+        self.px_per_nm_transformed = px_per_nm * self.resize_factor
+        
+        resized_image = resize_with_cv2(image, (int(h*self.resize_factor), int(w*self.resize_factor)))
+        
+        if w >= h:
+            pad_w = self.target_size-resized_image.size[0]
+            pad_h = max(0, self.target_short_edge_min-resized_image.size[1])
+        else:
+            pad_w = max(0, self.target_short_edge_min-resized_image.size[0])
+            pad_h = self.target_size-resized_image.size[1]
+        
+        # apply minimal padding
+        pad_w += 2*self.minimal_pad
+        pad_h += 2*self.minimal_pad
+        
+        # round to multiplication
+        pad_w += (self.pad_to_multiply - resized_image.size[0]%self.pad_to_multiply)%self.pad_to_multiply
+        pad_h += (self.pad_to_multiply - resized_image.size[1]%self.pad_to_multiply)%self.pad_to_multiply
+        
+        self.pad_right = pad_w // 2
+        self.pad_left = pad_w - self.pad_right
+        
+        self.pad_up = pad_h // 2
+        self.pad_bottom = pad_h - self.pad_up
+        
+        padded_image = tvf.pad(resized_image, [self.pad_left,self.pad_up, self.pad_right, self.pad_bottom], padding_mode='reflect') # fill 114 as in YOLO
+        return padded_image
+    
+    @property
+    def unpad_slice(self) -> Tuple[slice]:
+        return slice(self.pad_up,-self.pad_bottom if self.pad_bottom>0 else None), slice(self.pad_left,-self.pad_right if self.pad_right>0 else None)
+    
+    def inverse_transform(self, image: Union[np.ndarray, Image.Image], output_size: Optional[Tuple[int]]=None, output_nm_per_px: Optional[float]=None, return_pil: bool=True) -> Image.Image:
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        # h, w = image.shape[:2]
+        # unpadded_image = image[self.pad_up:h-self.pad_bottom,self.pad_left:w-self.pad_right]
+        unapdded_image = image[self.unpad_slice]
+        
+        if output_size is not None and output_nm_per_px is not None:
+            raise ValueError("one of output_size or output_nm_per_px must not be None")
+        elif output_nm_per_px is not None:
+            resize_factor = self.original_nm_per_px/output_nm_per_px
+            output_size = (int(self.original_size[0]*resize_factor), int(self.original_size[1]*resize_factor))
+        elif output_size is None:
+            output_size = self.original_size
+        resized_image = resize_with_cv2(unapdded_image, (output_size[1],output_size[0]), return_pil=return_pil) #Image.fromarray(unpadded_image).resize(self.original_size)
+        
+        return resized_image
+    
+def close_contour(contour):
+    if not np.array_equal(contour[0], contour[-1]):
+        contour = np.vstack((contour, contour[0]))
+    return contour
+
+def binary_mask_to_polygon(binary_mask, tol=0.01):
+    padded_binary_mask = np.pad(binary_mask, pad_width=1, mode='constant', constant_values=0)
+    contours = measure.find_contours(padded_binary_mask, 0.5)
+    # assert len(contours) == 1 #raise error if there are more than 1 contour
+    contour = contours[0]
+    contour -= 1 # correct for padding
+    contour = close_contour(contour)
+
+    polygon = measure.approximate_polygon(contour, tol)
+
+    polygon = np.flip(polygon, axis=1)
+    # after padding and subtracting 1 we may get -0.5 points in our polygon. Replace it with 0
+    polygon = np.where(polygon>=0, polygon, 0)
+    # segmentation = polygon.ravel().tolist()
+
+    return polygon
+
+def measure_shape(binary_mask):
+    contour = binary_mask_to_polygon(binary_mask)
+    perimeter = np.sum(np.linalg.norm(contour[:-1] - contour[1:], axis=1))
+    area = np.sum(binary_mask)
+        
+    return perimeter, area
+
+def calculate_gsi(perimeter, height, area):
+    a = 0.5*(perimeter - 2*height)
+    return 1 - area/(a*height)
+
+def object_slice(mask):
+    rows = np.any(mask, axis=1)
+    cols = np.any(mask, axis=0)
+    row_min, row_max = np.where(rows)[0][[0, -1]]
+    col_min, col_max = np.where(cols)[0][[0, -1]]
+
+    # Create a slice object for the bounding box
+    bounding_box_slice = (slice(row_min, row_max + 1), slice(col_min, col_max + 1))
+    
+    return bounding_box_slice
+
+
+def figure_to_pil(fig):
+    buf = BytesIO()
+    fig.savefig(buf, format='png')
+    buf.seek(0)
+
+    # Load the image from the buffer as a PIL Image
+    image = deepcopy(Image.open(buf))
+
+    # Close the buffer
+    buf.close()
+    return image
+
+def resize_to(image: Image.Image, s: int=4032, return_factor: bool =False) -> Image.Image:
+    w, h = image.size
+    longest_size = max(h, w)
+        
+    resize_factor = longest_size / s
+        
+    resized_image = image.resize((int(w/resize_factor), int(h/resize_factor)))
+    if return_factor:
+        return resized_image, resize_factor
+    return resized_image
+
+def resize_with_cv2(image, shape, return_pil=True):
+    """resize using cv2 with cv2.INTER_LINEAR - consistent with YOLO"""
+    h, w = shape
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+        
+    resized = cv2.resize(image, (w, h), interpolation=cv2.INTER_LINEAR)
+    if return_pil:
+        return Image.fromarray(resized)
+    else:
+        return resized
+
+def select_unique_mask(mask):
+    """if mask consists of multiple parts, select the largest"""
+    if not np.any(mask): # if mask is empty, return without change
+        return mask
+    blobs = ndimage.label(mask)[0]
+    blob_labels, blob_sizes = np.unique(blobs, return_counts=True)
+    best_blob_label = blob_labels[1:][np.argmax(blob_sizes[1:])]
+    return blobs == best_blob_label
+
+def sliced_mean(x, slice_size):
+    cs_y = np.cumsum(x, axis=0)
+    cs_y = np.concatenate((np.zeros((1, cs_y.shape[1]), dtype=cs_y.dtype), cs_y), axis=0)
+    slices_y = (cs_y[slice_size:] - cs_y[:-slice_size])/slice_size
+    cs_xy = np.cumsum(slices_y, axis=1)
+    cs_xy = np.concatenate((np.zeros((cs_xy.shape[0], 1), dtype=cs_xy.dtype), cs_xy), axis=1)
+    slices_xy = (cs_xy[:,slice_size:] - cs_xy[:,:-slice_size])/slice_size
+    return slices_xy
+
+def sliced_var(x, slice_size):
+    x = x.astype('float64')
+    return sliced_mean(x**2, slice_size) - sliced_mean(x, slice_size)**2
+
+def calculate_distance_map(mask):
+    padded = np.pad(mask, pad_width=1, mode='constant', constant_values=False)
+    distance_map_padded = ndimage.distance_transform_edt(padded)
+    return distance_map_padded[1:-1,1:-1]
+
+def select_samples(granum_image, granum_mask, crop_size=96, n_samples=64, granum_fraction_min=0.75, variance_p=0.):
+    granum_occupancy = sliced_mean(granum_mask, crop_size)
+    possible_indices = np.stack(np.where(granum_occupancy >= granum_fraction_min), axis=1)
+    
+    if variance_p == 0:
+        p = np.ones(len(possible_indices))
+    else:
+        variance_map = sliced_var(granum_image, crop_size)
+        p = variance_map[possible_indices[:,0], possible_indices[:,1]]**variance_p
+    p /= np.sum(p)
+    
+    chosen_indices = np.random.choice(
+        np.arange(len(possible_indices)),
+        min(len(possible_indices), n_samples),
+        replace=False,
+        p = p
+    )
+
+    crops = []
+    for crop_idx, idx in enumerate(chosen_indices):
+        crops.append(
+            granum_image[
+                possible_indices[idx,0]:possible_indices[idx,0]+crop_size,
+                possible_indices[idx,1]:possible_indices[idx,1]+crop_size
+            ]
+        )
+    return np.array(crops)
+
+def calculate_height(mask_oriented): #HACK
+    span = mask_oriented.shape[0] - np.argmax(mask_oriented[::-1], axis=0) - np.argmax(mask_oriented, axis=0)
+    return np.quantile(span, 0.8) 
+
+def calculate_diameter(mask_oriented):
+    """returns mean diameter"""
+    # calculate 0.25 and 0.75 lines
+    vertical_mask = np.any(mask_oriented, axis=1)
+    upper_granum_bound = np.argmax(vertical_mask)
+    lower_granum_bound = mask_oriented.shape[0] - np.argmax(vertical_mask[::-1])
+    upper = round(0.75*upper_granum_bound + 0.25*lower_granum_bound)
+    lower = max(upper+1, round(0.25*upper_granum_bound + 0.75*lower_granum_bound))
+    valid_rows_slice = slice(upper, lower)
+    
+    # calculate diameters
+    span = mask_oriented.shape[1] - np.argmax(mask_oriented[valid_rows_slice,::-1], axis=1) - np.argmax(mask_oriented[valid_rows_slice], axis=1)
+    return np.mean(span)
+
+def robust_mean(x, q=0.1):
+    x_med = np.median(x)
+    deviations = abs(x- x_med)
+    if max(deviations) == 0:
+        mask = np.ones(len(x), dtype='bool')
+    else:
+        threshold = np.quantile(deviations, 1-q)
+        mask = x[deviations<= threshold]
+     
+    return np.mean(x[mask])
+
+def rotate_image_and_mask(image, mask, direction):
+    mask_oriented = ndimage.rotate(mask.astype('int'), -direction, reshape=True).astype('bool')
+    idx_begin_x, idx_end_x = np.where(np.any(mask_oriented, axis=0))[0][np.array([0, -1])]
+    idx_begin_y, idx_end_y = np.where(np.any(mask_oriented, axis=1))[0][np.array([0, -1])]
+    img_oriented = ndimage.rotate(image, -direction, reshape=True) #[idx_begin_y:idx_end_y, idx_begin_x:idx_end_x]
+    return img_oriented, mask_oriented
+
+class GranaAnalyser:
+    def __init__(self, weights_detector: str, weights_orientation: str, weights_period: str, period_sd_threshold_nm: float=2.5) -> None:
+        """
+        Initializes the GranaAnalyser with specified weights for detection and measuring.
+
+        This method loads the weights for the grana detection and measuring algorithms 
+        from the specified file paths. It also loads mock data for visualization and 
+        analysis purposes.
+
+        Parameters:
+        weights_detector (str): The file path to the weights file for the grana detection algorithm.
+        weights_orientation (str): The file path to the weights file for the grana orientation algorithm.
+        weights_period (str): The file path to the weights file for the grana period algorithm.
+        """
+        self.detector = YOLO(weights_detector)
+        
+        self.orienter = PatchedPredictor(
+            StripsModelLumenWidth.load_from_checkpoint(weights_orientation, map_location='cpu').eval(),
+            normalization = dict(mean=0.250, std=0.135),
+            n_samples=32,
+            mask=None,
+            crop_size=64,
+            angle_confidence_threshold=0.2
+        )
+
+        self.measurement_px_per_nm = 1/0.768 # image scale required for measurement
+
+        self.period_measurer = PeriodMeasurer(
+            weights_period,
+            px_per_nm=self.measurement_px_per_nm,
+            sd_threshold_nm=period_sd_threshold_nm,
+            period_threshold_nm_min=14, period_threshold_nm_max=30
+        )
+        
+
+    def get_grana_data(self, image, detections, scaler, border_margin=1, min_count=1) -> List[Granum]:
+        """filter detections and create grana data"""
+        image_numpy = np.array(image)
+        if image_numpy.ndim == 3:
+            image_numpy = image_numpy[:,:,0]
+        
+        mask_all = None
+        grana = []
+        for mask, confidence in zip(
+            detections.masks.data.cpu().numpy().astype('bool'),
+            detections.boxes.conf.cpu().numpy()
+        ):
+            granum_mask = select_unique_mask(mask[scaler.unpad_slice])
+            # check if mask is empty after padding
+            if not np.any(granum_mask):
+                continue
+            granum_mask = ndimage.binary_fill_holes(granum_mask)
+
+            # check if touches boundary:
+            if (np.sum(granum_mask[:border_margin])>min_count) or \
+                (np.sum(granum_mask[-border_margin:])>min_count) or \
+                (np.sum(granum_mask[:,:border_margin])>min_count) or \
+                (np.sum(granum_mask[:,-border_margin:])>min_count):
+                
+                continue
+            
+            # check grana overlap
+            if mask_all is None:
+                mask_all = granum_mask
+            else:
+                intersection = mask_all & granum_mask
+                
+                if intersection.sum() >= (granum_mask.sum() * 0.2):
+                    continue
+                mask_all = mask_all | granum_mask
+            
+            granum = Granum(
+                image = image,
+                mask = granum_mask,
+                scaler=scaler,
+                detection_confidence=float(confidence)
+            )
+            
+            granum.image_numpy = image_numpy
+            grana.append(granum)
+        return grana
+    
+    def measure_grana(self, grana: List[Granum], measurement_image: np.ndarray) -> List[Granum]:
+        """measure grana: includes orientation detection, period detection and geometric measurements"""
+        for granum in grana:
+            measurement_mask = resize_with_cv2(granum.mask.astype(np.uint8), measurement_image.shape[:2], return_pil=False).astype('bool')
+            
+            granum.bounding_box_slice = object_slice(measurement_mask)
+            granum.image_crop = measurement_image[granum.bounding_box_slice][:,:]
+            granum.mask_crop = measurement_mask[granum.bounding_box_slice]
+                        
+            # initialize measurements
+            granum.measurements = {}
+
+            # measure shape
+            granum.measurements['perimeter px'], granum.measurements['area px'] = measure_shape(granum.mask_crop)
+
+            # measrure orientation
+            orienter_predictions = self.orienter(granum.image_crop, granum.mask_crop)
+            granum.measurements['direction'] = orienter_predictions["est_angle"]
+            granum.measurements['direction confidence'] = orienter_predictions["est_angle_confidence"]
+            
+            if not np.isnan(granum.measurements["direction"]):
+                img_oriented, mask_oriented = rotate_image_and_mask(granum.image_crop, granum.mask_crop, granum.measurements["direction"])
+                oriented_granum_slice = object_slice(mask_oriented)
+                granum.img_oriented = img_oriented[oriented_granum_slice]
+                granum.mask_oriented = mask_oriented[oriented_granum_slice]
+                granum.measurements['height px'] = calculate_height(granum.mask_oriented)
+                granum.measurements['GSI'] = calculate_gsi(
+                    granum.measurements['perimeter px'],
+                    granum.measurements['height px'],
+                    granum.measurements['area px']
+                )
+                granum.measurements['diameter px'] = calculate_diameter(granum.mask_oriented)
+                                
+                oriented_granum_slice = object_slice(granum.mask_oriented)
+                granum.measurements["period nm"], granum.measurements["period SD nm"] = self.period_measurer(granum.img_oriented, granum.mask_oriented)
+                
+                if not pd.isna(granum.measurements['period nm']):
+                    granum.measurements['Number of layers'] = round(granum.measurements['height px']/ self.measurement_px_per_nm / granum.measurements['period nm'])
+            
+        return grana
+
+    def extract_grana_data(self, grana: List[Granum]) -> pd.DataFrame:
+        """collect and scale grana data"""
+        grana_data = []
+        for granum in grana:
+            granum_entry = {
+                'Granum ID': granum.id,
+                'detection confidence': granum.detection_confidence
+            }
+            # fill with None if absent:
+            for key in ['direction', 'Number of layers', 'GSI', 'period nm', 'period SD nm']:
+                granum_entry[key] = granum.measurements.get(key, None)
+            # scale linearly:
+            for key in ['height px', 'diameter px', 'perimeter px', 'perimeter px']:
+                granum_entry[f"{key[:-3]} nm"] = granum.measurements.get(key, np.nan) / self.measurement_px_per_nm
+            # scale quadratically
+            granum_entry['area nm^2'] = granum.measurements['area px'] / self.measurement_px_per_nm**2
+            
+            grana_data.append(granum_entry)
+                
+        return pd.DataFrame(grana_data)
+    
+    def visualize_detections(self, grana: List[Granum], image: Image.Image) -> Image.Image:
+        visualization_longer_edge = 1024
+        scale = visualization_longer_edge/max(image.size)
+        visualization_size = (round(scale*image.size[0]), round(scale*image.size[1]))
+        visualization_image = np.array(image.resize(visualization_size).convert('RGB'))
+        
+        if len(grana) > 0:
+            grana_mask = resize_with_cv2(
+                np.any(np.array([granum.mask for granum in grana]),axis=0).astype(np.uint8),
+                visualization_size[::-1],
+                return_pil=False
+            ).astype('bool')
+            visualization_image[grana_mask]= (0.7*visualization_image[grana_mask] + 0.3*np.array([[[39, 179, 115]]])).astype(np.uint8)
+
+        
+        for granum in grana:
+            scale = visualization_longer_edge/max(granum.mask.shape)
+            y, x = ndimage.center_of_mass(granum.mask)
+            cv2.putText(visualization_image, f'{granum.id}', org=(int(x*scale)-10, int(y*scale)+10), fontFace=cv2.FONT_HERSHEY_SIMPLEX , fontScale=1, color=(39, 179, 115),thickness = 2)
+        
+        return Image.fromarray(visualization_image)       
+        
+
+    def generate_grana_images(self, grana: List[Granum], image_name: str ="") -> List[Image.Image]:
+        grana_images = {}
+        for granum in grana:
+            fig, ax = plt.subplots()
+            if granum.img_oriented is None:
+                image_to_plot = granum.image_crop
+                mask_to_plot = granum.mask_crop
+                extra_caption = " orientation and period unknown"
+            else:
+                image_to_plot = granum.img_oriented
+                mask_to_plot = granum.mask_oriented
+                extra_caption = ""
+            
+            ax.imshow(0.5*255*(~mask_to_plot) +image_to_plot*(1-0.5*(~mask_to_plot)), cmap='gray', vmin=0, vmax=255)
+            ax.axis('off')
+            ax.set_title(f'[{granum.id}]{image_name}\n{extra_caption}')
+            granum_image = figure_to_pil(fig)
+            grana_images[granum.id] = granum_image
+            plt.close('all')
+        
+        return grana_images
+
+    def format_data(self, grana_data: pd.DataFrame) -> pd.DataFrame:
+        rounding_roles = {'area nm^2': 0, 'perimeter nm': 1, 'diameter nm': 1, 'height nm': 1, 'period nm': 1, 'period SD nm': 2, 'GSI':2, 'direction': 1}
+        rounded_data = grana_data.round(rounding_roles)
+        columns_order = ['Granum ID', 'File name', 'area nm^2', 'perimeter nm', 'GSI','diameter nm', 'height nm', 'Number of layers','period nm', 'period SD nm', 'direction']
+        return rounded_data[columns_order]
+        
+    
+    def aggregate_data(self, grana_data: pd.DataFrame, confidence: Optional[float]=None) -> Dict:
+        if confidence is None:
+            filtered = grana_data
+        else:
+            filtered = grana_data.loc[grana_data['aggregated confidence'] >=  confidence]
+        aggregation = filtered[['area nm^2', 'perimeter nm', 'diameter nm', 'height nm', 'Number of layers', 'period nm', 'GSI']].mean().to_dict()
+        aggregation_std = filtered[['area nm^2', 'perimeter nm', 'diameter nm', 'height nm', 'Number of layers', 'period nm', 'GSI']].std().to_dict()
+        aggregation_std = {f"{k} std": v for k, v in aggregation_std.items()}
+        aggregation_result = {**aggregation, **aggregation_std, 'N grana': len(filtered)}
+        return aggregation_result
+            
+    def predict_on_single(self, image: Image.Image, scale: float, detection_confidence: float=0.25, granum_id_start=1, image_name: str = "") -> Tuple[List[Image.Image], pd.DataFrame, List[Image.Image]]:
+        """
+        Predicts and aggregates data related to grana using a dictionary of images.
+
+        Parameters:
+        image (Image.Image): PIL Image object to be analyzed
+        scale (float): scale of the image: px per nm.
+        detection_confidence (float): The detection confidence threshold shape measurement
+
+        Returns:
+        Tuple[Image.Image, pandas.DataFrame, List[Image.Image]]: 
+        A tuple containing:
+               - detection_visualization (Image.Image): PIL image representing 
+                 the detection visualizations.
+               - grana_data (pandas.DataFrame): A DataFrame containing the simulated granum data.
+               - grana_images (List[Image.Image]): A list of PIL images of the grana.
+        """
+        # convert to grayscale
+        image = image.convert("L")
+                
+        # detect
+        scaler = ScalerPadder(target_size=1024, target_short_edge_min=640)
+        scaled_image = scaler.transform(image, px_per_nm=scale)
+        detections = self.detector.predict(source=scaled_image, conf=detection_confidence)[0]
+
+        # get grana data
+        grana = self.get_grana_data(image, detections, scaler)
+        for granum_id, granum in enumerate(grana, start=granum_id_start):
+            granum.id = granum_id
+
+        
+        # visualize detections
+        detection_visualization = self.visualize_detections(grana, image)
+        
+        # measure grana
+        measurement_image_resize_factor = self.measurement_px_per_nm / scale
+        measurement_image_shape = (
+            int(image.size[1]*measurement_image_resize_factor),
+            int(image.size[0]*measurement_image_resize_factor)
+        )
+        measurement_image = resize_with_cv2(  # numpy image in scale valid for measurement
+            image, measurement_image_shape, return_pil=False
+        )
+        grana = self.measure_grana(grana, measurement_image)
+
+        # pandas DataFrame
+        grana_data = self.extract_grana_data(grana)
+        
+        # list of PIL images
+        grana_images = self.generate_grana_images(grana, image_name=image_name)
+        
+        return detection_visualization, grana_data, grana_images
+
+    def predict(self, images: Dict[str, Image.Image], scale: float, detection_confidence: float=0.25, parameter_confidence: Optional[float]=None) -> Tuple[List[Image.Image], pd.DataFrame, List[Image.Image], Dict]:
+        """
+        Predicts and aggregates data related to grana using a dictionary of images.
+
+        Parameters:
+        images (Dict[str, Image.Image]): A dictionary of PIL Image objects to be analyzed, 
+                                         keyed by their names.
+        scale (float): scale of the image: px per nm
+        detection_confidence (float): The detection confidence threshold shape measurement
+        parameter_confidence (float): The confidence threshold used for data aggregation. Only 
+                            data with aggregated confidence above this threshold will 
+                            be considered.
+
+        Returns:
+        Tuple[List[Image.Image], pandas.DataFrame, List[Image.Image], Dict]: 
+        A tuple containing:
+               - detection_visualizations (List[Image.Image]): A list of PIL images representing 
+                 the detection visualizations.
+               - grana_data (pandas.DataFrame): A DataFrame containing the simulated granum data.
+               - grana_images (List[Image.Image]): A list of PIL images of the grana.
+               - aggregated_data (Dict): A dictionary containing the aggregated data results.
+        """
+        detection_visualizations_all = {}
+        grana_data_all = None
+        grana_images_all = {}
+        
+        granum_id_start = 1
+        for image_name, image in images.items():
+            detection_visualization, grana_data, grana_images = self.predict_on_single(image, scale=scale, detection_confidence=detection_confidence, granum_id_start=granum_id_start, image_name=image_name)
+            granum_id_start += len(grana_data)
+            detection_visualizations_all[image_name] = detection_visualization
+            grana_images_all.update(grana_images)
+            
+            grana_data['File name'] = image_name
+            if grana_data_all is None:
+                grana_data_all = grana_data
+            else:
+                # grana_data['Granum ID'] += len(grana_data_all)
+                grana_data_all = pd.concat([grana_data_all, grana_data])
+            
+        # dict
+        # grana_data_all.to_csv('grana_data_all.csv', index=False)
+        aggregated_data = self.aggregate_data(grana_data_all, parameter_confidence)
+        
+        formatted_grana_data = self.format_data(grana_data_all)
+        
+        return detection_visualizations_all, formatted_grana_data, grana_images_all, aggregated_data
+
+    
+class GranaDetector(GranaAnalyser):
+    """supplementary class for grana detection only
+    """
+    def __init__(self, weights_detector: str, detector_config: Optional[str] = None, model_type="yolo") -> None:
+
+        if model_type == "yolo":
+            self.detector = YOLO(weights_detector)
+        elif model_type == "mmdetection":
+            self.detector = MMDetector(model=detector_config, weights=weights_detector)
+        else:
+            raise NotImplementedError()
+        
+    def predict_on_single(self, image: Image.Image, scale: float, detection_confidence: float=0.25, granum_id_start=1, use_scaling=True, granum_border_margin=1, granum_border_min_count=1, scaler_sizes=(1024, 640)) -> List[Granum]:
+        # convert to grayscale
+        image = image.convert("L")
+                
+        # detect
+        if use_scaling:
+            scaler = ScalerPadder(target_size=scaler_sizes[0], target_short_edge_min=scaler_sizes[1])
+        else:
+            #dummy scaler
+            scaler = ScalerPadder(target_size=max(image.size), target_short_edge_min=min(image.size), minimal_pad=0, pad_to_multiply=1)
+        scaled_image = scaler.transform(image, scale=scale)
+        detections = self.detector.predict(source=scaled_image, conf=detection_confidence)[0]
+
+        # get grana data
+        grana = self.get_grana_data(image, detections, scaler, border_margin=granum_border_margin, min_count=granum_border_min_count)
+        for i_granum, granum in enumerate(grana, start=1):
+            granum.id = i_granum
+        
+        return grana

period_calculation/config.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+transforms = [
+    A.Normalize(**{'mean': 0.2845, 'std': 0.1447}, max_pixel_value=1.0),
+    # Applies the formula (img - mean * max_pixel_value) / (std * max_pixel_value)
+    ToTensorV2()
+]
+
+model_config = {
+    'receptive_field_height': 220,
+    'receptive_field_width': 38,
+    'stride_height': 64,
+    'stride_width': 2,
+    'image_height': 476,
+    'image_width': 476}
+

period_calculation/data_reader.py

@@ -0,0 +1,861 @@
+import pandas as pd
+import numpy as np
+import skimage.io
+from pathlib import Path
+import torch
+import scipy
+from PIL import Image, ImageFilter, ImageChops
+# from config import model_config
+from period_calculation.config import model_config
+
+# Function to add Gaussian noise
+
+def add_microscope_noise(base_image_as_numpy, noise_intensity):
+    ###### The code below is for adding noise to the image
+    # noise intensity is a number between 0 and 1
+    # --- priginal implementation was provided by Michał Bykowski
+    # --- and adapted
+    # This routine works with PIL images and numpy internally (changing formats as it goes)
+    # but the input and output are numpy arrays
+
+    def add_noise(image, mean=0, std_dev=50):  # std_dev impacts the amount of noise
+        # Generating noise
+        noise = np.random.normal(mean, std_dev, (image.height, image.width))
+        # Adding noise to the image
+        noisy_image = np.array(image) + noise
+        # Ensuring the values remain within valid grayscale range
+        noisy_image = np.clip(noisy_image, 0, 255)
+        return Image.fromarray(noisy_image.astype('uint8'))
+
+
+    base_image = Image.fromarray(base_image_as_numpy)
+    gray_value = 128
+    gray = Image.new('L', base_image.size, color=gray_value)
+
+
+    gray = add_noise(gray, std_dev=noise_intensity * 76)
+    gray = gray.filter(ImageFilter.GaussianBlur(radius=3))
+    gray = add_noise(gray, std_dev=noise_intensity * 23)
+    gray = gray.filter(ImageFilter.GaussianBlur(radius=2))
+    gray = add_noise(gray, std_dev=noise_intensity * 15)
+
+    # soft light works as in Photoshop
+    # Superimposes two images on top of each other using the Soft Light algorithm
+    result = ImageChops.soft_light(base_image, gray)
+
+    return np.array(result)
+
+def detect_boundaries(mask, axis):
+    # calculate the boundaries of the mask
+    #axis = 0 results in x_from, x_to
+    #axis = 1 results in y_from, y_to
+
+
+    sum = mask.sum(axis=axis)
+
+    ind_from = min(sum.nonzero()[0])
+    ind_to = max(sum.nonzero()[0])
+    return ind_from, ind_to
+
+def add_symmetric_filling_beyond_mask(img, mask):
+    for x in range(img.shape[1]):
+        if sum(mask[:, x]) != 0:   #if there is at least one nonzero index
+            nonzero_indices = mask[:, x].nonzero()[0]
+
+            y_min = min(nonzero_indices)
+            y_max = max(nonzero_indices)
+
+            if y_max == y_min:   #there is only one point
+                img[:, x] = img[y_min, x]
+            else:
+                next = y_min + 1
+                step = +1  # we start by going upwards
+                for y in reversed(range(y_min)):
+                    img[y, x] = img[next, x]
+                    if next == y_max or next == y_min:   #we hit the boundaries - we reverse
+                        step *= -1   #reverse direction
+                    next += step
+
+                next = y_max - 1
+                step = -1  # we start by going downwards
+                for y in range(y_max + 1, img.shape[0]): #we hit the boundaries - we reverse
+                    img[y, x] = img[next, x]
+                    if next == y_max or next == y_min:
+                        step *= -1  # reverse direction
+                    next += step
+    return img
+class AbstractDataset(torch.utils.data.Dataset):
+
+    def __init__(self,
+                 model = None,
+                 transforms=[],
+                 #### distortions during training ####
+                 hv_symmetry=True,  # True or False
+
+                 min_horizontal_subsampling = 50,  # None to turn off; or minimal percentage of horizontal size of the image
+                 min_vertical_subsampling = 70,  # None to turn off; or minimal percentage of vertical size of the image
+                 max_random_tilt = 3,  # None to turn off; or maximum tilt in degrees
+                 max_add_colors_to_histogram = 10,  # 0 to turn off; or points of the histogram to be added
+                 max_remove_colors_from_histogram = 30,  # 0 to turn off; or points of the histogram to be removed
+                 max_noise_intensity = 3.0,  # 0.0 to turn off; or max intensity of the noise
+
+                 gaussian_phase_transforms_epoch=None,  # None to turn off; or number of the epoch when the gaussian phase starts
+                 min_horizontal_subsampling_gaussian_phase = 30,  # None to turn off; or minimal percentage of horizontal size of the image
+                 min_vertical_subsampling_gaussian_phase = 70,  # None to turn off; or minimal percentage of vertical size of the image
+                 max_random_tilt_gaussian_phase = 2,  # None to turn off; or maximum tilt in degrees
+                 max_add_colors_to_histogram_gaussian_phase = 10,  # 0 to turn off; or points of the histogram to be added
+                 max_remove_colors_from_histogram_gaussian_phase = 60,  # 0 to turn off; or points of the histogram to be removed
+                 max_noise_intensity_gaussian_phase = 3.5,  # 0.0 to turn off; or max intensity of the noise
+
+                 #### controling variables ####
+                 transform_level=2,   # 0 - no transforms, 1 - only the basic transform, 2 - all transforms, -1 - subsampling for high images
+                 retain_raw_images=False,
+                 retain_masks=False):
+
+
+        self.model = model  # we need that to check epoch number during training
+
+        self.hv_symmetry = hv_symmetry
+
+        self.min_horizontal_subsampling = min_horizontal_subsampling
+        self.min_vertical_subsampling = min_vertical_subsampling
+        self.max_random_tilt = max_random_tilt
+        self.max_add_colors_to_histogram = max_add_colors_to_histogram
+        self.max_remove_colors_from_histogram = max_remove_colors_from_histogram
+        self.max_noise_intensity = max_noise_intensity
+
+        self.gaussian_phase_transforms_epoch = gaussian_phase_transforms_epoch
+        self.min_horizontal_subsampling_gaussian_phase = min_horizontal_subsampling_gaussian_phase
+        self.min_vertical_subsampling_gaussian_phase = min_vertical_subsampling_gaussian_phase
+        self.max_random_tilt_gaussian_phase = max_random_tilt_gaussian_phase
+        self.max_add_colors_to_histogram_gaussian_phase = max_add_colors_to_histogram_gaussian_phase
+        self.max_remove_colors_from_histogram_gaussian_phase = max_remove_colors_from_histogram_gaussian_phase
+        self.max_noise_intensity_gaussian_phase = max_noise_intensity_gaussian_phase
+
+        self.image_height = model_config['image_height']
+        self.image_width = model_config['image_width']
+
+        self.transform_level = transform_level
+        self.retain_raw_images = retain_raw_images
+        self.retain_masks = retain_masks
+        self.transforms = transforms
+
+
+    def get_image_and_mask(self, row):
+        raise NotImplementedError("Subclass needs to implement this method")
+
+    def load_and_transform_image_and_mask(self, row):
+        img, mask = self.get_image_and_mask(row)
+
+        angle = row['angle']
+        #check if gaussian phase is on
+        if self.gaussian_phase_transforms_epoch is not None and self.model.current_epoch >= self.gaussian_phase_transforms_epoch:
+            max_random_tilt = self.max_random_tilt_gaussian_phase
+            max_noise_intensity = self.max_noise_intensity_gaussian_phase
+            min_horizontal_subsampling = self.min_horizontal_subsampling_gaussian_phase
+            min_vertical_subsampling = self.min_vertical_subsampling_gaussian_phase
+            max_add_colors_to_histogram = self.max_add_colors_to_histogram_gaussian_phase
+            max_remove_colors_from_histogram = self.max_remove_colors_from_histogram_gaussian_phase
+        else:
+            max_random_tilt = self.max_random_tilt
+            max_noise_intensity = self.max_noise_intensity
+            min_horizontal_subsampling = self.min_horizontal_subsampling
+            min_vertical_subsampling = self.min_vertical_subsampling
+            max_add_colors_to_histogram = self.max_add_colors_to_histogram
+            max_remove_colors_from_histogram = self.max_remove_colors_from_histogram
+
+
+
+
+
+
+        if self.transform_level >= 2 and max_random_tilt is not None:
+            ####### RANDOM TILT
+            angle += np.random.uniform(-max_random_tilt, max_random_tilt)
+
+        img = scipy.ndimage.rotate(img, 90 - angle, reshape=True, order=3)  # HORIZONTAL POSITION
+        ###the part of the image that is added after rotation is all black (0s)
+        mask = scipy.ndimage.rotate(mask, 90 - angle, reshape=True, order = 0)  # HORIZONTAL POSITION
+                    #order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+
+        ############# CROP
+        x_from, x_to = detect_boundaries(mask, axis=0)
+        y_from, y_to = detect_boundaries(mask, axis=1)
+
+        #crop the image to the verical and horizontal limits.
+        img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+        mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+
+        img_raw = img.copy()
+
+
+        if self.transform_level >= 2:
+            ########## ADDING NOISE
+
+            if max_noise_intensity > 0.0:
+                noise_intensity = np.random.random() * max_noise_intensity
+                noisy_img = add_microscope_noise(img, noise_intensity=noise_intensity)
+                img[mask] = noisy_img[mask]
+
+        if self.transform_level == -1:
+            #special case where we take at most 300 middle pixels from the image
+            # (vertical subsampling)
+            # to handle very latge images correctly
+            x_from, x_to = detect_boundaries(mask, axis=0)
+            y_from, y_to = detect_boundaries(mask, axis=1)
+
+            y_size = y_to - y_from + 1
+
+            random_size = 300   #not so random, ay?
+
+            if y_size > random_size:
+                random_start = y_size // 2 - random_size // 2
+
+                y_from = random_start
+                y_to = random_start + random_size - 1
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+                # recrop the image if necessary
+                # -- even after only horizontal subsampling it may be necessary to recrop the image
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+        if self.transform_level >= 1:
+            ############## HORIZONTAL SUBSAMPLING
+            if min_horizontal_subsampling is not None:
+                x_size = x_to - x_from + 1
+
+                # add some random horizontal shift
+                random_size = np.random.randint(x_size * min_horizontal_subsampling / 100.0, x_size + 1)
+                random_start = np.random.randint(0, x_size - random_size + 1) + x_from
+
+                img = img[:, random_start:(random_start + random_size)]
+                mask = mask[:, random_start:(random_start + random_size)]
+
+            ############ VERTICAL SUBSAMPLING
+            if min_vertical_subsampling is not None:
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                y_size = y_to - y_from + 1
+
+                random_size = np.random.randint(y_size * min_vertical_subsampling / 100.0, y_size + 1)
+                random_start = np.random.randint(0, y_size - random_size + 1) + y_from
+
+                y_from = random_start
+                y_to = random_start + random_size - 1
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+            if min_horizontal_subsampling is not None or min_vertical_subsampling is not None:
+                #recrop the image if necessary
+                # -- even after only horizontal subsampling it may be necessary to recrop the image
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+
+        ######### ADD SYMMETRIC FILLING OF THE IMAGE BEYOND THE MASK
+        #img = add_symmetric_filling_beyond_mask(img, mask)
+        #This leaves holes in the image, so we will not use it
+
+        #plt.imshow(img)
+        #plt.show()
+        ######### HORIZONTAL AND VERTICAL SYMMETRY.
+        # When superimposed, the result is 180 degree rotation
+        if self.transform_level >= 1 and self.hv_symmetry:
+            for axis in range(2):
+                if np.random.randint(0, 2) % 2 == 0:
+                    img = np.flip(img, axis = axis)
+                    mask = np.flip(mask, axis = axis)
+            #plt.imshow(img)
+            #plt.show()
+
+        if self.transform_level >= 2 and (max_add_colors_to_histogram > 0 or max_remove_colors_from_histogram > 0):
+            lower_bound = np.random.randint(-max_add_colors_to_histogram, max_remove_colors_from_histogram + 1)
+            upper_bound = np.random.randint(255 - max_remove_colors_from_histogram, 255 + max_add_colors_to_histogram + 1)
+            # first clip the values outstanding from the range (lower_bound -- upper_bound)
+            img[mask] = np.clip(img[mask], lower_bound, upper_bound)
+            # the range (lower_bound -- upper_bound) gets mapped to the range (0--255)
+            # but only in a portion of the image where mask = True
+            img[mask] = np.interp(img[mask], (lower_bound, upper_bound), (0, 255)).astype(np.uint8)
+
+        #### since preserve_range in skimage.transform.resize is set to False, the image
+        #### will be converted to float. Consult:
+        # https://scikit-image.org/docs/stable/api/skimage.transform.html#skimage.transform.resize
+        # https://scikit-image.org/docs/dev/user_guide/data_types.html
+
+        # In our case the image gets conveted to floats ranging 0-1
+        old_height = img.shape[0]
+        img = skimage.transform.resize(img, (self.image_height, self.image_width), order=3)
+        new_height = img.shape[0]
+        mask = skimage.transform.resize(mask, (self.image_height, self.image_width), order=0, preserve_range=True)
+           # order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+
+        scale_factor = new_height / old_height
+
+
+        #plt.imshow(img)
+        #plt.show()
+        #plt.imshow(mask)
+        #plt.show()
+        return img, mask, scale_factor, img_raw
+
+    def get_annotations_row(self, idx):
+        raise NotImplementedError("Subclass needs to implement this method")
+
+    def __getitem__(self, idx):
+        row = self.get_annotations_row(idx)
+
+        image, mask, scale_factor, image_raw = self.load_and_transform_image_and_mask(row)
+
+        image_data = {
+            'image': image,
+        }
+
+        for transform in self.transforms:
+            image_data = transform(**image_data)
+            # transform operates on image field ONLY of image_data, and returns a dictionary with the same keys
+
+        ret_dict = {
+            'image': image_data['image'],
+            'period_px': torch.tensor(row['period_nm'] * scale_factor * row['px_per_nm'], dtype=torch.float32),
+            'filename': row['granum_image'],
+            'px_per_nm': row['px_per_nm'],
+            'scale': scale_factor,         # the scale factor is used to calculate the true period error
+                                          # (before scale) in losses and metrics
+            'neutral': -self.transforms[0].mean/self.transforms[0].std   #value of 0 after the scale transform
+        }
+
+        if self.retain_raw_images:
+            ret_dict['image_raw'] = image_raw
+
+        if self.retain_masks:
+            ret_dict['mask'] = mask
+
+        return ret_dict
+
+    def __len__(self):
+        raise NotImplementedError("Subclass needs to implement this method")
+
+class ImageDataset(AbstractDataset):
+    def __init__(self, annotations, data_dir: Path, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.data_dir = Path(data_dir)
+
+        self.id = 1
+
+        if isinstance(annotations, str):
+            annotations = data_dir / annotations  #make it a Path object relative to data_dir
+
+        if isinstance(annotations, Path):
+            self.annotations = pd.read_csv(data_dir / annotations)
+            no_period = ['27_k7 [1]_4.png']
+            del_img = ['38_k42[1]_19.png', 'n6363_araLL_60kx_6 [1]_0.png', '27_hs8 [1]_5.png', '27_k7 [1]_20.png',
+                       'F1_1_60kx_01 [1]_2.png']
+            self.annotations = self.annotations[~self.annotations['granum_image'].isin(no_period)]
+            self.annotations = self.annotations[~self.annotations['granum_image'].isin(del_img)]
+        else:
+            self.annotations = annotations
+
+    def get_image_and_mask(self, row):
+        filename = row['granum_image']
+        img_path = self.data_dir / filename
+        img_raw = skimage.io.imread(img_path)
+
+        img = img_raw[:, :, 0]   # all three channels are equal, with the exception
+        # of the last channel which is the full blue (0,0,255) for outside the mask (so blue channel is 255, red and green are 0)
+        mask = (img_raw != (0, 0, 255)).any(axis=2)
+        return img, mask
+
+    def get_annotations_row(self, idx):
+        row = self.annotations.iloc[idx].to_dict()
+        row['idx'] = idx
+        return row
+
+    def __len__(self):
+        return len(self.annotations)
+
+class ArtificialDataset(AbstractDataset):
+    def __init__(self,
+                 min_period = 20,
+                 max_period = 140,
+                 white_fraction_min = 0.15,
+                 white_fraction_max=0.45,
+
+                 noise_min_sd = 0.0,
+                 noise_max_sd = 100.0,
+                 noise_max_sd_everywhere = 20.0,  # 20.0
+                 leftovers_max = 5,
+
+                 get_real_masks_dataset = None,   #None or instance of ImageDataset
+                 *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.id = 0
+        self.min_period = min_period
+        self.max_period = max_period
+        self.white_fraction_min = white_fraction_min
+        self.white_fraction_max = white_fraction_max
+
+        self.receptive_field_height = model_config['receptive_field_height']
+        self.stride_height = model_config['stride_height']
+        self.receptive_field_width = model_config['receptive_field_width']
+        self.stride_width = model_config['stride_width']
+
+        self.noise_min_sd = noise_min_sd
+        self.noise_max_sd = noise_max_sd
+        self.noise_max_sd_everywhere = noise_max_sd_everywhere
+
+        self.leftovers_max = leftovers_max
+
+        self.get_real_masks_dataset = get_real_masks_dataset
+
+
+    def get_image_and_mask(self, row):
+        # generate a rectangular image of black and white horizontal stripes
+        # with black stripes varying with white stripes
+
+        period_px = row['period_nm'] * row['px_per_nm']
+        # white occupying 5-20 % of a total period (white+black)
+        white_px = np.random.randint(period_px * self.white_fraction_min, period_px * self.white_fraction_max + 1)
+
+
+        # mask is rectangle of True values
+        img = np.zeros((self.image_height, self.image_width), dtype=np.uint8)
+        mask = np.ones((self.image_height, self.image_width), dtype=bool)
+        black_px = period_px - white_px
+        random_start = np.random.randint(0, period_px+1)
+        for i in range(self.image_height):
+            if (random_start+i) % (black_px + white_px) < black_px:
+                # sample width with random numbers from 0 to 101
+                img[i, :] = np.random.randint(0, 101, self.image_width)
+            else:
+                #sample width with random numbers from 156 to 255
+                img[i, :] = np.random.randint(156, 256, self.image_width)
+
+        if self.noise_max_sd_everywhere > self.noise_min_sd:
+            sd = np.random.uniform(self.noise_min_sd, self.noise_max_sd_everywhere)
+            noise = np.random.normal(0, sd, (self.image_height, self.image_width))
+            img = np.clip(img+noise.astype(img.dtype), 0, 255)
+
+        if self.noise_max_sd > self.noise_min_sd:
+            # there is also a metagrid in the image
+            # consisting of overlapping receptive fields of size 190x42
+            # with stride 64x4
+            # the metagrid is 5x102
+            overlapping_fields_count_height = (self.image_height - self.receptive_field_height) // self.stride_height + 1
+            overlapping_fields_count_width = (self.image_width - self.receptive_field_width) // self.stride_width + 1
+
+
+            sd = np.random.uniform(self.noise_min_sd, self.noise_max_sd)
+            noise = np.random.normal(0, sd, (self.image_height, self.image_width))
+
+            #there will be some left-over metagrid rectangles
+            leftovers_count = np.random.randint(1, self.leftovers_max)
+            for i in range(leftovers_count):
+                metagrid_row = np.random.randint(0, overlapping_fields_count_height)
+                metagrid_col = np.random.randint(0, overlapping_fields_count_width)
+                #zero-out the noise inside the selected metagrid
+                noise[metagrid_row * self.stride_height:metagrid_row * self.stride_height + self.receptive_field_height + 1, \
+                      metagrid_col * self.stride_width :metagrid_col * self.stride_width + self.receptive_field_width + 1] = 0
+
+            #add noise to the image
+            img = np.clip(img+noise.astype(img.dtype), 0, 255)
+
+            if self.get_real_masks_dataset is not None:
+                ret_dict = self.get_real_masks_dataset.__getitem__(row['idx'] % len(self.get_real_masks_dataset))
+                mask = ret_dict['mask']   #this mask is already sized target height-by-width
+
+                img[mask == False] = 0
+
+        return img, mask
+
+    def get_annotations_row(self, idx):
+        return {'idx': idx,
+                'period_nm': np.random.randint(self.min_period, self.max_period),
+                'px_per_nm': 1.0,
+                'granum_image': 'artificial_%d.png' % idx,
+                'angle': 90}
+
+
+    def __len__(self):
+        return 237 # number of samples as in real data in the train set (70% of 339 is 237,3)
+
+    
+class AdHocDataset(AbstractDataset):
+    def __init__(self, images_masks_pxpernm: list[tuple[np.ndarray, np.ndarray, float]], *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.data = images_masks_pxpernm
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        image, mask, px_per_nm = self.data[idx]
+
+        image, mask, scale_factor, image_raw = self.load_and_transform_image_and_mask(image, mask)
+
+        image_data = {
+            'image': image,
+        }
+
+        for transform in self.transforms:
+            image_data = transform(**image_data)
+            # transform operates on image field ONLY of image_data, and returns a dictionary with the same keys
+
+        ret_dict = {
+            'image': image_data['image'],
+            'period_px': torch.tensor(0, dtype=torch.float32),
+            'filename': str(idx),
+            'px_per_nm': px_per_nm,
+            'scale': scale_factor,         # the scale factor is used to calculate the true period error
+                                          # (before scale) in losses and metrics
+            'neutral': -self.transforms[0].mean/self.transforms[0].std   #value of 0 after the scale transform
+        }
+
+        if self.retain_raw_images:
+            ret_dict['image_raw'] = image_raw
+
+        if self.retain_masks:
+            ret_dict['mask'] = mask
+
+        return ret_dict
+        
+        
+    def load_and_transform_image_and_mask(self, img, mask):
+        
+        angle = 90
+        #check if gaussian phase is on
+        if self.gaussian_phase_transforms_epoch is not None and self.model.current_epoch >= self.gaussian_phase_transforms_epoch:
+            max_random_tilt = self.max_random_tilt_gaussian_phase
+            max_noise_intensity = self.max_noise_intensity_gaussian_phase
+            min_horizontal_subsampling = self.min_horizontal_subsampling_gaussian_phase
+            min_vertical_subsampling = self.min_vertical_subsampling_gaussian_phase
+            max_add_colors_to_histogram = self.max_add_colors_to_histogram_gaussian_phase
+            max_remove_colors_from_histogram = self.max_remove_colors_from_histogram_gaussian_phase
+        else:
+            max_random_tilt = self.max_random_tilt
+            max_noise_intensity = self.max_noise_intensity
+            min_horizontal_subsampling = self.min_horizontal_subsampling
+            min_vertical_subsampling = self.min_vertical_subsampling
+            max_add_colors_to_histogram = self.max_add_colors_to_histogram
+            max_remove_colors_from_histogram = self.max_remove_colors_from_histogram
+
+
+        if self.transform_level >= 2 and max_random_tilt is not None:
+            ####### RANDOM TILT
+            angle += np.random.uniform(-max_random_tilt, max_random_tilt)
+        
+
+        img = scipy.ndimage.rotate(img, 90 - angle, reshape=True, order=3)  # HORIZONTAL POSITION
+        ###the part of the image that is added after rotation is all black (0s)
+        mask = scipy.ndimage.rotate(mask, 90 - angle, reshape=True, order = 0)  # HORIZONTAL POSITION
+                    #order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+
+        ############# CROP
+        x_from, x_to = detect_boundaries(mask, axis=0)
+        y_from, y_to = detect_boundaries(mask, axis=1)
+
+        #crop the image to the verical and horizontal limits.
+        img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+        mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+
+        img_raw = img.copy()
+
+
+        if self.transform_level >= 2:
+            ########## ADDING NOISE
+
+            if max_noise_intensity > 0.0:
+                noise_intensity = np.random.random() * max_noise_intensity
+                noisy_img = add_microscope_noise(img, noise_intensity=noise_intensity)
+                img[mask] = noisy_img[mask]
+
+        if self.transform_level == -1:
+            #special case where we take at most 300 middle pixels from the image
+            # (vertical subsampling)
+            # to handle very latge images correctly
+            x_from, x_to = detect_boundaries(mask, axis=0)
+            y_from, y_to = detect_boundaries(mask, axis=1)
+
+            y_size = y_to - y_from + 1
+
+            random_size = 300   #not so random, ay?
+
+            if y_size > random_size:
+                random_start = y_size // 2 - random_size // 2
+
+                y_from = random_start
+                y_to = random_start + random_size - 1
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+                # recrop the image if necessary
+                # -- even after only horizontal subsampling it may be necessary to recrop the image
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+        if self.transform_level >= 1:
+            ############## HORIZONTAL SUBSAMPLING
+            if min_horizontal_subsampling is not None:
+                x_size = x_to - x_from + 1
+
+                # add some random horizontal shift
+                random_size = np.random.randint(x_size * min_horizontal_subsampling / 100.0, x_size + 1)
+                random_start = np.random.randint(0, x_size - random_size + 1) + x_from
+
+                img = img[:, random_start:(random_start + random_size)]
+                mask = mask[:, random_start:(random_start + random_size)]
+
+            ############ VERTICAL SUBSAMPLING
+            if min_vertical_subsampling is not None:
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                y_size = y_to - y_from + 1
+
+                random_size = np.random.randint(y_size * min_vertical_subsampling / 100.0, y_size + 1)
+                random_start = np.random.randint(0, y_size - random_size + 1) + y_from
+
+                y_from = random_start
+                y_to = random_start + random_size - 1
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+            if min_horizontal_subsampling is not None or min_vertical_subsampling is not None:
+                #recrop the image if necessary
+                # -- even after only horizontal subsampling it may be necessary to recrop the image
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+
+        ######### ADD SYMMETRIC FILLING OF THE IMAGE BEYOND THE MASK
+        #img = add_symmetric_filling_beyond_mask(img, mask)
+        #This leaves holes in the image, so we will not use it
+
+        #plt.imshow(img)
+        #plt.show()
+        ######### HORIZONTAL AND VERTICAL SYMMETRY.
+        # When superimposed, the result is 180 degree rotation
+        if self.transform_level >= 1 and self.hv_symmetry:
+            for axis in range(2):
+                if np.random.randint(0, 2) % 2 == 0:
+                    img = np.flip(img, axis = axis)
+                    mask = np.flip(mask, axis = axis)
+            #plt.imshow(img)
+            #plt.show()
+
+        if self.transform_level >= 2 and (max_add_colors_to_histogram > 0 or max_remove_colors_from_histogram > 0):
+            lower_bound = np.random.randint(-max_add_colors_to_histogram, max_remove_colors_from_histogram + 1)
+            upper_bound = np.random.randint(255 - max_remove_colors_from_histogram, 255 + max_add_colors_to_histogram + 1)
+            # first clip the values outstanding from the range (lower_bound -- upper_bound)
+            img[mask] = np.clip(img[mask], lower_bound, upper_bound)
+            # the range (lower_bound -- upper_bound) gets mapped to the range (0--255)
+            # but only in a portion of the image where mask = True
+            img[mask] = np.interp(img[mask], (lower_bound, upper_bound), (0, 255)).astype(np.uint8)
+
+        #### since preserve_range in skimage.transform.resize is set to False, the image
+        #### will be converted to float. Consult:
+        # https://scikit-image.org/docs/stable/api/skimage.transform.html#skimage.transform.resize
+        # https://scikit-image.org/docs/dev/user_guide/data_types.html
+
+        # In our case the image gets conveted to floats ranging 0-1
+        old_height = img.shape[0]
+        img = skimage.transform.resize(img, (self.image_height, self.image_width), order=3)
+        new_height = img.shape[0]
+        mask = skimage.transform.resize(mask, (self.image_height, self.image_width), order=0, preserve_range=True)
+           # order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+
+        scale_factor = new_height / old_height
+
+
+        #plt.imshow(img)
+        #plt.show()
+        #plt.imshow(mask)
+        #plt.show()
+        return img, mask, scale_factor, img_raw
+
+    
+class AdHocDataset2(AbstractDataset):
+    def __init__(self, images_masks_pxpernm: list[tuple[np.ndarray, np.ndarray, float]], *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.data = images_masks_pxpernm
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        image, mask, px_per_nm = self.data[idx]
+
+        image, mask, scale_factor, image_raw = self.load_and_transform_image_and_mask(image, mask)
+
+        image_data = {
+            'image': image,
+        }
+
+        for transform in self.transforms:
+            image_data = transform(**image_data)
+            # transform operates on image field ONLY of image_data, and returns a dictionary with the same keys
+
+        ret_dict = {
+            'image': image_data['image'],
+            'scale': scale_factor,         # the scale factor is used to calculate the true period error
+                                          # (before scale) in losses and metrics
+            'neutral': -self.transforms[0].mean/self.transforms[0].std   #value of 0 after the scale transform
+        }
+
+        return ret_dict
+        
+        
+    def load_and_transform_image_and_mask(self, img, mask):
+ 
+        img_raw = img.copy()
+
+        if self.transform_level == -1:
+            #special case where we take at most 300 middle pixels from the image
+            # (vertical subsampling)
+            # to handle very latge images correctly
+            x_from, x_to = detect_boundaries(mask, axis=0)
+            y_from, y_to = detect_boundaries(mask, axis=1)
+
+            y_size = y_to - y_from + 1
+
+            max_size = 300
+
+            if y_size > max_size:
+                random_start = y_size // 2 - max_size // 2
+
+                y_from = random_start
+                y_to = random_start + max_size - 1
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+                # recrop the image if necessary
+                # -- even after only horizontal subsampling it may be necessary to recrop the image
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+
+        #### since preserve_range in skimage.transform.resize is set to False, the image
+        #### will be converted to float. Consult:
+        # https://scikit-image.org/docs/stable/api/skimage.transform.html#skimage.transform.resize
+        # https://scikit-image.org/docs/dev/user_guide/data_types.html
+
+        # In our case the image gets conveted to floats ranging 0-1
+        old_height = img.shape[0]
+        img = skimage.transform.resize(img, (self.image_height, self.image_width), order=3)
+        new_height = img.shape[0]
+        mask = skimage.transform.resize(mask, (self.image_height, self.image_width), order=0, preserve_range=True)
+           # order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+
+        scale_factor = new_height / old_height
+
+        return img, mask, scale_factor, img_raw
+
+class AdHocDataset3(AbstractDataset):
+    def __init__(self, images_and_masks: list[tuple[np.ndarray, np.ndarray]], *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.data = images_and_masks
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        image, mask = self.data[idx]
+
+        image, mask, scale_factor = self.load_and_transform_image_and_mask(image, mask)
+
+        image_data = {
+            'image': image,
+        }
+
+        for transform in self.transforms:
+            image_data = transform(**image_data)
+            # transform operates on image field ONLY of image_data, and returns a dictionary with the same keys
+
+        ret_dict = {
+            'image': image_data['image'],
+            'scale': scale_factor,         # the scale factor is used to calculate the true period error
+                                          # (before scale) in losses and metrics
+               #value of 0 after the scale transform
+        }
+
+        return ret_dict
+        
+        
+    def load_and_transform_image_and_mask(self, img, mask):
+ 
+        if self.transform_level == -1:
+            #special case where we take at most 300 middle pixels from the image
+            # (vertical subsampling)
+            # to handle very latge images correctly
+            x_from, x_to = detect_boundaries(mask, axis=0)
+            y_from, y_to = detect_boundaries(mask, axis=1)
+
+            y_size = y_to - y_from + 1
+
+            max_size = 300
+
+            if y_size > max_size:
+                random_start = y_size // 2 - max_size // 2
+
+                y_from = random_start
+                y_to = random_start + max_size - 1
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+                # recrop the image if necessary
+                # -- even after only horizontal subsampling it may be necessary to recrop the image
+
+                x_from, x_to = detect_boundaries(mask, axis=0)
+                y_from, y_to = detect_boundaries(mask, axis=1)
+
+                img = img[y_from:(y_to + 1), x_from:(x_to + 1)]
+                mask = mask[y_from:(y_to + 1), x_from:(x_to + 1)]
+
+
+        #### since preserve_range in skimage.transform.resize is set to False, the image
+        #### will be converted to float. Consult:
+        # https://scikit-image.org/docs/stable/api/skimage.transform.html#skimage.transform.resize
+        # https://scikit-image.org/docs/dev/user_guide/data_types.html
+
+        # In our case the image gets conveted to floats ranging 0-1
+        old_height = img.shape[0]
+        img = skimage.transform.resize(img, (self.image_height, self.image_width), order=3)
+        new_height = img.shape[0]
+        mask = skimage.transform.resize(mask, (self.image_height, self.image_width), order=0, preserve_range=True)
+           # order = 0 is the nearest neighbor interpolation, so the mask is not interpolated
+
+        scale_factor = new_height / old_height
+
+        return img, mask, scale_factor

period_calculation/image_transforms.py

@@ -0,0 +1,79 @@
+import numpy as np
+import torch
+import cv2
+
+
+def np_batched_radon(image_batch):
+    #image_batch:    torch tensor, batch_size x 1 x img_size x img_size
+    # squeeze order #1 and transform to numpy
+
+    image_batch = image_batch.squeeze(1).cpu().numpy()
+
+    batch_size, img_size = image_batch.shape[:2]
+    if batch_size > 512: # limit batch size to 512 because cv2.warpAffine fails for batch> 512
+        return np.concatenate([np_batched_radon(image_batch[i:i+512]) for i in range(0,batch_size,512)], axis=0)
+    theta = np.arange(180)
+    radon_image = np.zeros((image_batch.shape[0], img_size, len(theta)),
+                           dtype='float32')
+
+    for i, angle in enumerate(theta):
+        M = cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1)
+        rotated = cv2.warpAffine(np.transpose(image_batch, (1, 2, 0)),M,(img_size,img_size))
+
+        #plt.imshow(rotated[:,:,0])
+        #plt.show()
+
+        if batch_size == 1: # cv2.warpAffine cancels batch dimension if equal to 1
+          rotated = rotated[:,:, np.newaxis]
+        rotated = np.transpose(rotated, (2, 0, 1)) / 224.0
+        #rotated = rotated / np.array(255, dtype='float32')
+        radon_image[:, :, i] = rotated.sum(axis=1)
+
+    #plot the image
+
+  #  plt.imshow(radon_image[0])
+   # plt.show()
+
+    return radon_image
+
+
+def torch_batched_radon(image_batch, neutral_value):
+    #image_batch:                                batch_size x 1 x img_size x img_size
+    #np_batched_radon(image_batch - neutral_value)
+
+    image_batch = image_batch - neutral_value   # so the 0 value is neutral
+
+
+    batch_size = image_batch.shape[0]
+    img_size = image_batch.shape[2]
+
+    theta = np.arange(180)   # we don't need torch here, we will evaluate individual angles below
+
+    radon_image = torch.zeros((batch_size, 1, img_size, len(theta)), dtype=torch.float, device=image_batch.device)
+
+
+    for i, angle in enumerate(theta):
+        #M = cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1)
+        #calculate the same rotation matrix but with torch:
+        M = torch.tensor(cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1)).to(image_batch.device, dtype=torch.float32)
+        angle = torch.tensor((angle+90)/180.0*np.pi)
+        M1 = torch.tensor([[torch.sin(angle), torch.cos(angle), 0],
+                    [torch.cos(angle), -torch.sin(angle), 0]]).to(image_batch.device, dtype=torch.float32)
+
+
+        # we need to add a batch dimension to the rotation matrix
+        M1 = M1.repeat(batch_size, 1, 1)
+
+        grid = torch.nn.functional.affine_grid(M1, image_batch.shape, align_corners=False)
+        rotated = torch.nn.functional.grid_sample(image_batch, grid, mode='bilinear', padding_mode='zeros', align_corners=False)
+        rotated = rotated.squeeze(1)
+
+        #plt.imshow(rotated[0].cpu().numpy())
+        #plt.show()
+
+        radon_image[:, 0, :, i] = rotated.sum(axis=1) / 224.0 + neutral_value
+
+    #plt.imshow(radon_image[0, 0].cpu().numpy())
+    #plt.show()
+
+    return radon_image

period_calculation/models/abstract_model.py

@@ -0,0 +1,61 @@
+import pytorch_lightning as pl
+import torch
+
+from hydra.utils import instantiate
+
+
+class AbstractModel(pl.LightningModule):
+    def __init__(self,
+                 lr=0.001,
+                 optimizer_hparams=dict(),
+                 scheduler=dict(classname='MultiStepLR', kwargs=dict(milestones=[100, 150], gamma=0.1))
+                 ):
+        super().__init__()
+        # Exports the hyperparameters to a YAML file, and create "self.hparams" namespace
+        self.save_hyperparameters()
+
+    def forward(self, x):
+        raise NotImplementedError("Subclass needs to implement this method")
+
+    def configure_optimizers(self):
+        # AdamW is Adam with a correct implementation of weight decay (see here
+        # for details: https://arxiv.org/pdf/1711.05101.pdf)
+        print("configuring the optimizer and lr scheduler with learning rate=%.5f"%self.hparams.lr)
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr, **self.hparams.optimizer_hparams)
+        # scheduler = getattr(torch.optim.lr_scheduler, self.hparams.lr_hparams['classname'])(optimizer, **self.hparams.lr_hparams['kwargs'])
+        if self.hparams.scheduler is not None:
+            scheduler = instantiate({**self.hparams.scheduler, '_partial_': True})(optimizer)
+
+            return [optimizer], [scheduler]
+        else:
+            return optimizer
+
+    def additional_losses(self):
+        """get additional_losses"""
+        return torch.zeros((1))
+
+    def process_batch_supervised(self, batch):
+        """get predictions, losses and mean errors (MAE)"""
+        raise NotImplementedError("Subclass needs to implement this method")
+
+    def log_all(self, losses, metrics, prefix=''):
+        for k, v in losses.items():
+            self.log(f'{prefix}{k}_loss', v.item() if isinstance(v, torch.Tensor) else v)
+
+        for k, v in metrics.items():
+            self.log(f'{prefix}{k}', v.item() if isinstance(v, torch.Tensor) else v)
+
+    def training_step(self, batch, batch_idx):
+        # "batch" is the output of the training data loader.
+        preds, losses, metrics = self.process_batch_supervised(batch)
+        self.log_all(losses, metrics, prefix='train_')
+
+        return losses['final']
+
+    def validation_step(self, batch, batch_idx):
+        preds, losses, metrics = self.process_batch_supervised(batch)
+        self.log_all(losses, metrics, prefix='val_')
+
+    def test_step(self, batch, batch_idx):
+        preds, losses, metrics = self.process_batch_supervised(batch)
+        self.log_all(losses, metrics, prefix='test_')

period_calculation/models/gauss_model.py

@@ -0,0 +1,237 @@
+import pytorch_lightning as pl
+import torch
+import numpy as np
+
+from period_calculation.models.abstract_model import AbstractModel
+
+from period_calculation.config import model_config  # this is a dictionary with the model configuration
+
+class GaussPeriodModel(AbstractModel):
+    def __init__(self,
+                 *args, **kwargs
+                 ):
+        super().__init__(*args, **kwargs)
+
+        self.seq = torch.nn.Sequential(
+            torch.nn.Conv2d(1, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.MaxPool2d((2, 2), stride=(2, 2)),
+
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.MaxPool2d((2, 1), stride=(2, 1)),
+
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.MaxPool2d((2, 1), stride=(2, 1)),
+
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.MaxPool2d((2, 1), stride=(2, 1)),
+
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.MaxPool2d((2, 1), stride=(2, 1)),
+
+            torch.nn.Conv2d(32, 32, (3, 3), stride=(1, 1), padding=(0, 0)),
+            torch.nn.MaxPool2d((2, 1), stride=(2, 1)),
+
+            torch.nn.Dropout(0.1)
+        )
+        self.query = torch.nn.Parameter(torch.empty(1, 2, 32))   #two heads only
+        torch.nn.init.xavier_normal_(self.query)
+
+        self.linear1 = torch.nn.Linear(64, 8)
+        self.linear2 = torch.nn.Linear(8, 1)
+
+        self.query_sd = torch.nn.Parameter(torch.empty(1, 2, 32))
+        torch.nn.init.xavier_normal_(self.query_sd)
+
+        self.linear_sd1 = torch.nn.Linear(64, 8)
+        self.linear_sd2 = torch.nn.Linear(8, 1)
+        self.relu = torch.nn.ReLU()
+
+
+    def copy_network_trunk(self, model):
+        # https://discuss.pytorch.org/t/copy-weights-from-only-one-layer-of-one-model-to-another-model-with-different-structure/153419
+        with torch.no_grad():
+            for i, layer in enumerate(model.seq):
+                if i%2 == 0 and i!=20:   #convolutional layers are the ones with even indexes with the exeption of the 20th (=dropout)
+                    self.seq[i].weight.copy_(layer.weight)
+                    self.seq[i].bias.copy_(layer.bias)
+
+
+    def copy_final_layers(self, model):
+        # https://discuss.pytorch.org/t/copy-weights-from-only-one-layer-of-one-model-to-another-model-with-different-structure/153419
+
+        with torch.no_grad():
+            self.linear1.weight.copy_(model.linear1.weight)
+            self.linear1.bias.copy_(model.linear1.bias)
+
+            self.linear2.weight.copy_(model.linear2.weight)
+            self.linear2.bias.copy_(model.linear2.bias)
+
+            self.query.copy_(model.query)
+
+    def duplicate_final_layers(self):
+        # https://discuss.pytorch.org/t/copy-weights-from-only-one-layer-of-one-model-to-another-model-with-different-structure/153419
+
+        with torch.no_grad():
+            self.linear_sd1.weight.copy_(self.linear1.weight)
+            self.linear_sd1.bias.copy_(self.linear1.bias)
+
+            self.linear_sd2.weight.copy_(self.linear2.weight/10)
+            self.linear_sd2.bias.copy_(self.linear2.bias/10)
+
+            self.query_sd.copy_(self.query)
+
+    def forward(self, x, neutral=None, return_raw=False):
+        #https://www.nature.com/articles/s41598-023-43852-x
+
+        # x is sized                            # batch x 1 x 476 x 476
+
+        preds = self.seq(x)                        # batch x 32 x 5 x 220
+        features = torch.flatten(preds, 2)         # batch x 32 x 1100
+
+    # attention
+        energy = self.query @ features                           # batch x 2 x 1100
+        weights = torch.nn.functional.softmax(energy, 2)    # batch x 2 x 1100
+        response = features @ weights.transpose(1, 2) # batch x 32 x 2
+        response = torch.flatten(response, 1)         # batch x 64
+
+        preds = self.linear1(response)             # batch x 8
+        preds = self.linear2(self.relu(preds))     # batch x 1
+
+    # attention sd
+
+        energy_sd = self.query_sd @ features                         # batch x 2 x 1100
+        weights_sd = torch.nn.functional.softmax(energy_sd, 2)  # batch x 2 x 1100
+        response_sd = features @ weights_sd.transpose(1, 2)  # batch x 32 x 2
+        response_sd = torch.flatten(response_sd, 1)  # batch x 64
+
+        preds_sd = self.linear_sd1(response_sd)  # batch x 8
+        preds_sd = self.linear_sd2(self.relu(preds_sd))  # batch x 1
+
+        outputs = [ model_config['receptive_field_height']/(preds[:,0])  , torch.exp(preds_sd[:,0]) ]
+        if return_raw:
+            outputs.append(preds)
+            outputs.append(preds_sd)
+            outputs.append(weights)
+            outputs.append(weights_sd)
+
+        return tuple(outputs)
+
+    def additional_losses(self):
+        """get additional_losses"""
+        # additional (orthogonal) loss
+        # we multiply the two heads and later the MSE loss (towards zero) sums the result in L2 norm
+        # the idea is that the scalar product of two orthogonal vectors is zero
+        scalar_product = torch.cat((self.query[0, 0] * self.query[0, 1], self.query_sd[0, 0] * self.query_sd[0, 1]), dim=0)
+        orthogonal_loss = torch.nn.functional.mse_loss(scalar_product, torch.zeros_like(scalar_product))
+        return orthogonal_loss
+
+
+
+    def process_batch_supervised(self, batch):
+        """get predictions, losses and mean errors (metrics)"""
+
+        # get predictions
+        preds = {}
+        preds['period_px'], preds['sd'] = self.forward(batch['image'], batch['neutral'][0], return_raw=False)  # preds: period, sd, orto, preds_raw
+
+        # https://johaupt.github.io/blog/NN_prediction_uncertainty.html
+        # calculate losses
+        mse_period_px = torch.nn.functional.mse_loss(batch['period_px'],
+                                                     preds['period_px'])
+
+        gaussian_nll = torch.nn.functional.gaussian_nll_loss(batch['period_px'],
+                                                             preds['period_px'],
+                                                             (preds['sd']) ** 2)
+
+        orthogonal_weight = 0.1
+        orthogonal_loss = self.additional_losses()
+        length_of_the_first_phase = 0
+        if self.current_epoch < length_of_the_first_phase:
+            #transition from MSE to Gaussian Negative Log Likelihood with sin/cos over first epochs
+            angle = torch.tensor((self.current_epoch) / (length_of_the_first_phase) * np.pi / 2)
+            total_loss = (gaussian_nll) * torch.sin(angle) + (mse_period_px) * torch.cos(angle) + orthogonal_weight * orthogonal_loss
+        else:
+            total_loss = gaussian_nll + orthogonal_weight * orthogonal_loss
+
+        losses = {
+            'gaussian_nll': gaussian_nll,
+            'mse_period_px': mse_period_px,
+            'orthogonal': orthogonal_loss,
+            'final': total_loss
+        }
+
+        # calculate mean errors
+        ground_truth_detached = batch['period_px'].detach().cpu().numpy()
+        print(ground_truth_detached)
+        mean_detached = preds['period_px'].detach().cpu().numpy()
+        print(mean_detached)
+        sd_detached = preds['sd'].detach().cpu().numpy()
+        print("==>", sd_detached)
+        px_per_nm_detached = batch['px_per_nm'].detach().cpu().numpy()
+        scale_detached = batch['scale'].detach().cpu().numpy()
+
+        period_px_difference = np.mean(abs(
+            ground_truth_detached - mean_detached
+        ))
+
+        #initiate both with python array with 5 zeros
+        true_period_px_difference = [0.0] * 5
+        true_period_nm_difference = [0.0] * 5
+
+        for i, dist in enumerate([1.0, 2.0, 3.0, 4.0, 5.0]):
+            true_period_px_difference[i] = (np.sum(abs(
+                ((ground_truth_detached - mean_detached) / scale_detached) * (sd_detached / scale_detached <dist))) \
+                / np.sum(sd_detached / scale_detached < dist)) if np.sum(sd_detached / scale_detached < dist) > 0 else 0
+
+        for i, dist in enumerate([1.0, 2.0, 3.0, 4.0, 5.0]):
+            true_period_nm_difference[i] = (np.sum(abs(
+                ((ground_truth_detached - mean_detached) / (scale_detached * px_per_nm_detached)) * (sd_detached / scale_detached <dist))) \
+                / np.sum(sd_detached / scale_detached < dist)) if np.sum(sd_detached / scale_detached < dist) > 0 else 0
+
+
+        true_period_px_difference_all = np.mean(abs(
+            ((ground_truth_detached - mean_detached) / scale_detached)
+        ))
+
+        true_period_nm_difference_all = np.mean(abs(
+            ((ground_truth_detached - mean_detached) / (scale_detached * px_per_nm_detached))
+        ))
+
+        metrics = {
+            'period_px': period_px_difference,
+            'true_period_px_1': true_period_px_difference[0],
+            'true_period_px_2': true_period_px_difference[1],
+            'true_period_px_3': true_period_px_difference[2],
+            'true_period_px_4': true_period_px_difference[3],
+            'true_period_px_5': true_period_px_difference[4],
+            'true_period_px_all': true_period_px_difference_all,
+
+            'true_period_nm_1': true_period_nm_difference[0],
+            'true_period_nm_2': true_period_nm_difference[1],
+            'true_period_nm_3': true_period_nm_difference[2],
+            'true_period_nm_4': true_period_nm_difference[3],
+            'true_period_nm_5': true_period_nm_difference[4],
+            'true_period_nm_all': true_period_nm_difference_all,
+
+            'count_1': np.sum(sd_detached / scale_detached < 1.0),
+            'count_2': np.sum(sd_detached / scale_detached < 2.0),
+            'count_3': np.sum(sd_detached / scale_detached < 3.0),
+            'count_4': np.sum(sd_detached / scale_detached < 4.0),
+            'count_5': np.sum(sd_detached / scale_detached < 5.0),
+
+            'count_all': np.sum(sd_detached > 0.0),
+            'mean_sd': np.mean(sd_detached),
+            'sd_sd': np.std(sd_detached),
+        }
+
+        return preds, losses, metrics
+
+

period_calculation/period_measurer.py

@@ -0,0 +1,54 @@
+import torch
+from pathlib import Path
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import skimage
+import scipy
+import numpy as np
+from pytorch_lightning import seed_everything
+
+from period_calculation.data_reader import AdHocDataset3
+from period_calculation.models.gauss_model import GaussPeriodModel
+
+
+transforms = [
+    A.Normalize(**{'mean': 0.2845, 'std': 0.1447}, max_pixel_value=1.0),
+    # Applies the formula (img - mean * max_pixel_value) / (std * max_pixel_value)
+    ToTensorV2()
+]
+
+class PeriodMeasurer:
+    """returns period in pixels"""
+    def __init__(
+        self, weights_file, image_height=476, image_width=476,
+        px_per_nm = 1,
+        sd_threshold_nm=np.inf,
+        period_threshold_nm_min=0, period_threshold_nm_max=np.inf):
+        
+        self.model = GaussPeriodModel.load_from_checkpoint(weights_file).to("cpu") #.eval()?
+        self.px_per_nm = px_per_nm
+        self.sd_threshold_nm = sd_threshold_nm
+        self.period_threshold_nm_min = period_threshold_nm_min
+        self.period_threshold_nm_max = period_threshold_nm_max
+        
+    def __call__(self, img: np.ndarray, mask: np.ndarray) -> float:
+        seed_everything(44)
+        dataset = AdHocDataset3(
+            images_and_masks = [(img, mask)],
+            transform_level=-1,
+            retain_raw_images=False,
+            transforms=transforms
+        )
+        
+        image_data = dataset[0]
+        with torch.no_grad():
+            y_hat, sd_hat = self.model(image_data["image"].unsqueeze(0), return_raw=False)
+        
+        y_hat_nm = (y_hat/image_data["scale"]).item() / self.px_per_nm
+        sd_hat_nm = (sd_hat/image_data["scale"]).item() /self.px_per_nm
+        
+        
+        if (sd_hat_nm>self.sd_threshold_nm) or (y_hat_nm<self.period_threshold_nm_min) or (y_hat_nm>self.period_threshold_nm_max):
+            y_hat_nm = np.nan
+            
+        return y_hat_nm, sd_hat_nm

requirements.txt

@@ -0,0 +1,11 @@
+pandas==2.1.1
+torch==2.1.0
+torchvision==0.16
+ultralytics==8.0.216
+scikit-image==0.22.0
+pytorch-lightning==2.1.2
+timm==0.9.11
+albumentations==1.4.10
+hydra-core==1.3.2
+gradio==4.44.0
+albucore==0.0.16

settings.py

@@ -0,0 +1 @@
+DEMO = False

styles.css

@@ -0,0 +1,47 @@
+.header {
+    display: flex;
+    padding: 30px;
+    text-align: center;
+    justify-content: center;
+}
+
+#header-text {
+    font-size: 50px;
+    line-height: 50px;
+}
+
+#header-logo {
+    width: 50px;
+    height: 50px;
+    margin-right: 10px;
+    /*background-image: url("file=images/logo.svg");*/
+}
+
+.input-row {
+    max-width: 900px;
+    margin: 0 auto;
+}
+
+.margin-bottom {
+    margin-bottom: 48px;
+}
+
+.results-header {
+    margin-top: 48px;
+    text-align: center;
+    font-size: 45px;
+    margin-bottom: 12px;
+}
+
+.processed-info {
+    display: flex;
+    padding: 30px;
+    text-align: center;
+    justify-content: center;
+    font-size: 26px;
+}
+
+.title {
+    margin-bottom: 8px!important;
+    font-size: 22px;
+}

weights/AS_square_v16.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c1d7d4e56f0ea28b34dd2457807e9266d0d5539cf5adfb0719fed10791f79c5
+size 44771917

weights/model_weights_detector.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37fe3d98789572cf147d5f0d1ea99d50a57e5c2028454c3825295e89b6350fd5
+size 23926765

weights/period_measurer_weights-1.298_real_full-fa12970.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90133b01545ae9d30eafcbdec480410e0f8b9fae3ba1aabb280450ce9589100a
+size 350396

weights/yolo/20240604_yolov8_segm_ABRCR1_all_train4_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:038555e9b9f3900ec29c9c79d7b3d5b50aa3ca37b3e665ee7aa2394facc7e20e
+size 23926765

weights/yolo/current_yolo.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37fe3d98789572cf147d5f0d1ea99d50a57e5c2028454c3825295e89b6350fd5
+size 23926765