# CZMDOEL package and *.czann Model Exchange Format

This project provides simple-to-use conversion tools to generate a CZANN file from a [ONNX](https://onnx.ai/) model that resides in memory or on disk to be usable in the
[ZEN AI Toolkit](https://www.zeiss.com/microscopy/en/products/software/zeiss-zen/ai-toolkit.html) module starting with ZEN blue >=3.2 and ZEN Core >3.0.

In addition the models can be used in [ZEISS arivis platform](https://www.zeiss.com/microscopy/en/products/software/advanced-image-analysis.html)

Please check the following compatibility matrix for ZEN Blue/Core and the respective version (self.version) of the CZANN Model Specification JSON Meta data file (see _CZANN Model Specification_ below). Version compatibility is defined via the [Semantic Versioning Specification (SemVer)](https://semver.org/lang/de/).

| Model (legacy)/JSON | ZEN Blue | ZEN Core |
| ------------------- | :------: | -------: |
| 2.0.0               | \>= 3.13 | \>= 3.13 |
| 1.1.0               | \>= 3.5  |  \>= 3.4 |
| 1.0.0               | \>= 3.5  |  \>= 3.4 |
| 3.1.0 (legacy)      | \>= 3.4  |  \>= 3.3 |
| 3.0.0 (legacy)      | \>= 3.2  |  \>= 3.1 |

If you encounter a version mismatch when importing a model into ZEN, please check for the correct version of this package.

## Structure of repo

This repo is divided into two separate packages -: core, pytorch.

- Core - Provides base functionality, no dependency on Pytorch required.
- PyTorch - Provides PyTorch-specific functionalities, and converters based on PyTorch-logics.

## Installation

The library provides a base package and extras for export functionalities that require specific dependencies -:

- ```pip install czmodel``` - This would only install base dependencies, no Pytorch-specific packages.
- ```pip install czmodel[pytorch]``` - This would install base and Pytorch-specific packages.

## Samples

### For czmodel[pytorch]

- For single-class semantic segmentation:&nbsp;
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zeiss-microscopy/OAD/blob/master/Machine_Learning/notebooks/czmodel/SingleClassSemanticSegmentation_PyTorch_6_1_0.ipynb)

- For regression:&nbsp;
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zeiss-microscopy/OAD/blob/master/Machine_Learning/notebooks/czmodel/Regresssion_PyTorch_6_1_0.ipynb)

## System setup

The current version of this toolbox only requires a fresh Python >= 3.10, <3.14 installation.

## Model conversion

The toolbox provides a `convert` module that features all supported conversion strategies. It currently supports
converting PyTorch models in memory or stored on disk with a corresponding metadata JSON file (see _CZANN Model Specification_ below).

### PyTorch models in memory

The toolbox also provides functionality that can be imported e.g. in the training script used to fit a PyTorch model.
It provides different converters to target specific versions of the export format. Currently, there are two converters available:

- DefaultConverter: Exports a .czann file complying with the specification below.
- LegacyConverter (Only for segmentation): Exports a .czmodel file (version 3.1.0 of the legacy ANN-based segmentation models in ZEN).

The converters are accessible by running:

For PyTorch model:

```python
from czmodel.pytorch.convert import DefaultConverter, LegacyConverter
```

Every converter provides a `convert_from_model_spec` function that uses a model specification object to convert a model
to the corresponding export format. It accepts a `torch.nn.Module` that will be exported to [ONNX](https://onnx.ai/)
format and at the same time wrapped into a .czann/.czmodel file that can be imported and used by Intellesis.
To provide the metadata, the toolbox provides a ModelSpec class that must be filled with the model, a ModelMetadata
instance containing the information required by the specification (see _Model Metadata_ below), and optionally a license file.

A CZANN/CZMODEL can be created from a PyTorch model with the following three steps.

#### 1. Create a model metadata class

To export a CZANN, meta information is needed that must be provided through a `ModelMetadata` instance.

For segmentation:

```python
from czmodel.core.model_metadata import ModelMetadata, ModelType

model_metadata = ModelMetadata(
    input_shape=[1024, 1024, 3],
    output_shape=[1024, 1024, 5],
    model_type=ModelType.SINGLE_CLASS_SEMANTIC_SEGMENTATION,
    classes=["class1", "class2", "class3", "class4", "class5"],
    model_name="ModelName",
    min_overlap=[90, 90]
)
```

For regression:

```python
from czmodel.core.model_metadata import ModelMetadata, ModelType

model_metadata = ModelMetadata(
    input_shape=[1024, 1024, 3],
    output_shape=[1024, 1024, 3],
    model_type=ModelType.REGRESSION,
    model_name="ModelName",
    min_overlap=[90, 90]
)
```

For legacy CZMODEL models the legacy `ModelMetadata` must be used:

```python
from czmodel.core.legacy_model_metadata import ModelMetadata as LegacyModelMetadata

model_metadata_legacy = LegacyModelMetadata(
    name="Simple_Nuclei_SegmentationModel_Legacy",
    classes=["class1", "class2"],
    pixel_types="Bgr24",
    color_handling="ConvertToMonochrome",
    border_size=90,
)
```

#### 2 .Creating a model specification

The model and its corresponding metadata are now wrapped into a ModelSpec object.

```python
from czmodel.pytorch.model_spec import ModelSpec   # for czmodel[pytorch]

model_spec = ModelSpec(
    model=model,
    model_metadata=model_metadata,
    license_file="C:\\some\\path\\to\\a\\LICENSE.txt"
)
```

The corresponding model spec for legacy models is instantiated analogously.

```python
from czmodel.pytorch.legacy_model_spec import ModelSpec as LegacyModelSpec    # for czmodel[pytorch]

legacy_model_spec = LegacyModelSpec(
    model=model,
    model_metadata=model_metadata_legacy,
    license_file="C:\\some\\path\\to\\a\\LICENSE.txt"
)
```

#### 3. Converting the model

The actual model conversion is finally performed with the ModelSpec object and the output path and name of the CZANN.

```python
from czmodel.pytorch.convert import DefaultConverter as PytorchDefaultConverter  # for czmodel[pytorch]

PytorchDefaultConverter().convert_from_model_spec(model_spec=model_spec, output_path='some/path', output_name='some_file_name', input_shape=(3, 1024, 1024))
```

For legacy models the interface is similar.

```python
from czmodel.pytorch.convert import LegacyConverter as PytorchLegacyConverter  # for czmodel[pytorch]

PytorchLegacyConverter().convert_from_model_spec(model_spec=legacy_model_spec, output_path='some/path',
                                          output_name='some_file_name', input_shape=(3, 1024, 1024))
```

### Exported PyTorch models

Not all PyTorch models can be converted. You can convert a model exported from PyTorch if the model and the
provided metadata comply with the _CZANN Model Specification_ below.

The actual conversion is triggered by either calling:

```python
from czmodel.pytorch.convert import DefaultConverter as PytorchDefaultConverter  # for czmodel[pytorch]

PytorchDefaultConverter().convert_from_json_spec('Path to JSON file', 'Output path', (3, 1024, 1024), 'Model Name')
```

### Unpacking CZANN/CZSEG files

The czmodel library offers functionality to unpack existing CZANN/CZSEG models. For a given .czann or .czseg model it is possible to extract the underlying ANN model to a specified folder and retrieve the corresponding meta-data as instances of the meta-data classes defined in the czmodel library.

For CZANN files:

```python
from czmodel.pytorch.convert import DefaultConverter    # for czmodel[pytorch]
from pathlib import Path

model_metadata, model_path = DefaultConverter().unpack_model(model_file='Path of the .czann file',
                                                             target_dir=Path('Output Path'))
```

For CZSEG/CZMODEL files:

```python
from czmodel.pytorch.convert import DefaultConverter   # for czmodel[pytorch]
from pathlib import Path

model_metadata, model_path = LegacyConverter().unpack_model(model_file='Path of the .czseg/.czann file',
                                                            target_dir=Path('Output Path'))
```

## CZANN Model Specification

This section specifies the requirements for an artificial neural network (ANN) model and the additionally required metadata to enable execution of the model inside the ZEN AI Toolkit (Intellesis) infrastructure starting with ZEN blue >=3.2 and ZEN Core >3.0.

The model format currently allows to bundle models for semantic segmentation, instance segmentation, object detection, classification and regression and is defined as a ZIP archive with the file extension .czann containing the following files with the respective filenames:

- JSON Meta data file. (filename: model.json)
- Model in ONNX or PyTorch format. In case of  SavedModel format the folder representing the model must be zipped to a single file. (filename: model.model)
- Optionally: A license file for the contained model. (filename: license.txt)

The meta data file must comply with the following specification:

```json
{
    "$schema": "http://iglucentral.com/schemas/com.snowplowanalytics.self-desc/schema/jsonschema/1-0-0#",
    "$id": "http://127.0.0.1/model_format.schema.json",
    "title": "Exchange format for ANN models",
    "description": "A format that defines the meta information for exchanging ANN models.\nAny future versions of this specification should be evaluated through https://docs.snowplowanalytics.com/docs/pipeline-components-and-applications/iglu/igluctl-0-7-2/#lint-1 with --skip-checks numericMinMax,stringLength,optionalNull and https://www.json-buddy.com/json-schema-analyzer.htm.",
    "type": "object",
    "self": {
        "vendor": "com.zeiss",
        "name": "model-format",
        "format": "jsonschema",
        "version": "2-0-0"
    },
    "properties": {
        "Id": {
            "description": "Universally unique identifier of 128 bits for the model.",
            "type": "string"
        },
        "Type": {
            "description": "The type of problem addressed by the model.",
            "type": "string",
            "enum": [
                "SingleClassInstanceSegmentation",
                "MultiClassInstanceSegmentation",
                "SingleClassSemanticSegmentation",
                "MultiClassSemanticSegmentation",
                "SingleClassClassification",
                "MultiClassClassification",
                "ObjectDetection",
                "Regression"
            ]
        },
        "MinOverlap": {
            "description": "The minimum overlap of tiles for each dimension in pixels. Must be divisible by two.\nIn tiling strategies that consider tile borders instead of overlaps the minimum overlap is twice the border size.",
            "type": "array",
            "items": {
                "description": "The overlap of a single spatial dimension",
                "type": "integer",
                "minimum": 0
            },
            "minItems": 1
        },
        "Classes": {
            "description": "The class names corresponding to the last output dimension of the prediction.\nIf the last dimension of the prediction has shape n the provided list must be of length n",
            "type": "array",
            "items": {
                "description": "A name describing a class for segmentation and classification tasks",
                "type": "string"
            },
            "minItems": 2
        },
        "ModelName": {
            "description": "The name of exported neural network model in ONNX (file) or PyTorch format in the same ZIP archive as the meta data file.\nIn the case of ONNX the model must use ONNX opset version 12.\nhe model must contain exactly one input node which must comply with the input shape defined in the InputShape parameter\nand must have a batch dimension as its first dimension that is either 1 or undefined.",
            "type": "string"
        },
        "InputShape": {
            "description": "The shape of an input image. A typical 2D model has an input of shape [h, w, c] where h and w are the spatial dimensions and c is the number of channels.\nA 3D model is expected to have an input shape of [z, h, w, c] that contains an additional dimension z which represents the third spatial dimension.\nThe batch dimension is not specified here.\nThe input of the model must be of type float32 in the range [0..1].",
            "type": "array",
            "items": {
                "description": "The size of a single dimension",
                "type": "integer",
                "minimum": 1
            },
            "minItems": 3,
            "maxItems": 4
        },
        "OutputShape": {
            "description": "The shape of the output image. A typical 2D model has an input of shape [h, w, c] where h and w are the spatial dimensions and c is the number of classes.\nA 3D model is expected to have an input shape of [z, h, w, c] that contains an additional dimension z which represents the third spatial dimension.\nThe batch dimension is not specified here.\nIf the output of the model represents an image, it must be of type float32 in the range [0..1].",
            "type": "array",
            "items": {
                "description": "The size of a single dimension",
                "type": "integer",
                "minimum": 1
            },
            "minItems": 3,
            "maxItems": 4
        },
        "Scaling": {
            "description": "The extent of a pixel in x, y and z-direction (in that order) in units of m.",
            "type": "array",
            "items": {
                "description": "The extent of a pixel or voxel in a single dimension in units of m",
                "type": "number"
            },
            "minItems": 2,
            "maxItems": 3
        }
    },
    "required": [
        "Id",
        "Type",
        "InputShape",
        "OutputShape"
    ]
}
```

Json files can contain escape sequences and \\-characters in paths must be escaped with \\\\.

The following code snippet shows an example for a valid metadata file:

For single-class semantic segmentation:

```json
{
  "Id": "b511d295-91ff-46ca-bb60-b2e26c393809",
  "Type": "SingleClassSemanticSegmentation",
  "Classes": ["class1", "class2", "class3", "class4", "class5"],
  "InputShape": [1024, 1024, 3],
  "OutputShape": [1024, 1024, 5]
}
```

For regression:

```json
{
  "Id": "064587eb-d5a1-4434-82fc-2fbc9f5871f9",
  "Type": "Regression",
  "InputShape": [1024, 1024, 3],
  "OutputShape": [1024, 1024, 3]
}
```
