careamics 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

careamics/dataset_ng/demos/demo_datamodule.ipynb

@@ -0,0 +1,447 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from careamics.config.configuration_factories import (\n",
+    "    _create_ng_data_configuration,\n",
+    "    _list_spatial_augmentations,\n",
+    ")\n",
+    "from careamics.lightning.dataset_ng.data_module import CareamicsDataModule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import tifffile\n",
+    "from careamics_portfolio import PortfolioManager\n",
+    "\n",
+    "# instantiate data portfolio manage\n",
+    "portfolio = PortfolioManager()\n",
+    "\n",
+    "# and download the data\n",
+    "root_path = Path(\"./data\")\n",
+    "files = portfolio.denoising.N2V_BSD68.download(root_path)\n",
+    "\n",
+    "# create paths for the data\n",
+    "data_path = Path(root_path / \"denoising-N2V_BSD68.unzip/BSD68_reproducibility_data\")\n",
+    "train_path = data_path / \"train\"\n",
+    "val_path = data_path / \"val\"\n",
+    "test_path = data_path / \"test\" / \"images\"\n",
+    "gt_path = data_path / \"test\" / \"gt\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_std, image_mean = [], []\n",
+    "for file in train_path.glob(\"*.tiff\"):\n",
+    "    image = tifffile.imread(file)\n",
+    "    image_std.append(image.std())\n",
+    "    image_mean.append(image.mean())\n",
+    "image_std, image_mean = np.mean(image_std), np.mean(image_mean)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from path, train and val, no target\n",
+    "\n",
+    "config = _create_ng_data_configuration(\n",
+    "    data_type=\"tiff\",\n",
+    "    axes=\"SYX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=64,\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
+    "\n",
+    "config.set_means_and_stds([image_mean], [image_std])\n",
+    "config.val_dataloader_params = {\"shuffle\": False}\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    train_data=train_path,\n",
+    "    val_data=val_path,\n",
+    ")\n",
+    "data_module.setup('fit')\n",
+    "data_module.setup('validate')\n",
+    "\n",
+    "train_batch = next(iter(data_module.train_dataloader()))\n",
+    "val_batch = next(iter(data_module.val_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from path, only predict\n",
+    "\n",
+    "from careamics.config.data import NGDataConfig\n",
+    "\n",
+    "config = NGDataConfig(\n",
+    "    data_type=\"tiff\",\n",
+    "    patching={\n",
+    "        \"name\": \"tiled\",\n",
+    "        \"patch_size\": (128, 128),\n",
+    "        \"overlaps\": (32, 32)\n",
+    "    },\n",
+    "    axes=\"YX\",\n",
+    "    batch_size=8,\n",
+    "    image_means=[image_mean],\n",
+    "    image_stds=[image_std]\n",
+    ")\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    pred_data=test_path\n",
+    ")\n",
+    "data_module.setup('predict')\n",
+    "\n",
+    "pred_batch = next(iter(data_module.predict_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(pred_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# test from array"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_array = tifffile.imread(sorted(train_path.rglob('*'))[0])\n",
+    "val_array = tifffile.imread(sorted(val_path.rglob('*'))[0])\n",
+    "test_array = tifffile.imread(sorted(test_path.rglob('*'))[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from array, train and val, no target\n",
+    "\n",
+    "config = _create_ng_data_configuration(\n",
+    "    data_type=\"array\",\n",
+    "    axes=\"SYX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=64,\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
+    "\n",
+    "config.set_means_and_stds([image_mean], [image_std])\n",
+    "config.val_dataloader_params = {\"shuffle\": False}\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    train_data=train_array,\n",
+    "    val_data=val_array,\n",
+    ")\n",
+    "data_module.setup('fit')\n",
+    "data_module.setup('validate')\n",
+    "\n",
+    "train_batch = next(iter(data_module.train_dataloader()))\n",
+    "val_batch = next(iter(data_module.val_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# test with target"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import skimage\n",
+    "\n",
+    "example_data = skimage.data.human_mitosis()\n",
+    "\n",
+    "markers = np.zeros_like(example_data)\n",
+    "markers[example_data < 25] = 1\n",
+    "markers[example_data > 50] = 2\n",
+    "\n",
+    "elevation_map = skimage.filters.sobel(example_data)\n",
+    "segmentation = skimage.segmentation.watershed(elevation_map, markers)\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 2)\n",
+    "ax[0].imshow(example_data)\n",
+    "ax[1].imshow(segmentation)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = _create_ng_data_configuration(\n",
+    "    data_type=\"array\",\n",
+    "    axes=\"YX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=64,\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
+    "config.set_means_and_stds(\n",
+    "    [example_data.mean()],\n",
+    "    [example_data.std()],\n",
+    "    [segmentation.mean()],\n",
+    "    [segmentation.std()]\n",
+    ")\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    train_data=[example_data],\n",
+    "    train_data_target=[segmentation],\n",
+    "    val_data=[example_data],\n",
+    "    val_data_target=[segmentation]\n",
+    ")\n",
+    "data_module.setup('fit')\n",
+    "data_module.setup('validate')\n",
+    "\n",
+    "train_batch = next(iter(data_module.train_dataloader()))\n",
+    "val_batch = next(iter(data_module.val_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(2, 8, figsize=(10, 3))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[0][i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "    ax[1][i].imshow(train_batch[1].data[i][0].numpy())\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots(2, 8, figsize=(10, 3))\n",
+    "for i in range(8):\n",
+    "    ax[0][i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "    ax[1][i].imshow(val_batch[1].data[i][0].numpy())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from array, only predict, with target\n",
+    "\n",
+    "from careamics.config.data import NGDataConfig\n",
+    "\n",
+    "config = NGDataConfig(\n",
+    "    data_type=\"array\",\n",
+    "    patching={\n",
+    "        \"name\": \"tiled\",\n",
+    "        \"patch_size\": (128, 128),\n",
+    "        \"overlaps\": (32, 32)\n",
+    "    },\n",
+    "    axes=\"YX\",\n",
+    "    batch_size=8,\n",
+    "    image_means=[image_mean],\n",
+    "    image_stds=[image_std]\n",
+    ")\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    pred_data=example_data,\n",
+    "    pred_data_target=segmentation\n",
+    ")\n",
+    "data_module.setup('predict')\n",
+    "\n",
+    "pred_batch = next(iter(data_module.predict_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(pred_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from list of paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = _create_ng_data_configuration(\n",
+    "    data_type=\"tiff\",\n",
+    "    axes=\"SYX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=64,\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
+    "\n",
+    "config.set_means_and_stds([image_mean], [image_std])\n",
+    "config.val_dataloader_params = {\"shuffle\": False}\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    train_data=sorted(train_path.glob(\"*.tiff\")),\n",
+    "    val_data=sorted(val_path.glob(\"*.tiff\")),\n",
+    ")\n",
+    "data_module.setup('fit')\n",
+    "data_module.setup('validate')\n",
+    "\n",
+    "train_batch = next(iter(data_module.train_dataloader()))\n",
+    "val_batch = next(iter(data_module.val_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from custom"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = _create_ng_data_configuration(\n",
+    "    data_type=\"custom\",\n",
+    "    axes=\"SYX\",\n",
+    "    patch_size=(64, 64),\n",
+    "    batch_size=64,\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
+    "\n",
+    "config.set_means_and_stds([image_mean], [image_std])\n",
+    "config.val_dataloader_params = {\"shuffle\": False}\n",
+    "\n",
+    "def read_source_func(path):\n",
+    "    image = tifffile.imread(path)\n",
+    "    image = 255 - image\n",
+    "    return image\n",
+    "\n",
+    "data_module = CareamicsDataModule(\n",
+    "    data_config=config,\n",
+    "    train_data=sorted(train_path.glob(\"*.tiff\")),\n",
+    "    val_data=sorted(val_path.glob(\"*.tiff\")),\n",
+    "    read_source_func=read_source_func\n",
+    ")\n",
+    "data_module.setup('fit')\n",
+    "data_module.setup('validate')\n",
+    "\n",
+    "train_batch = next(iter(data_module.train_dataloader()))\n",
+    "val_batch = next(iter(data_module.val_dataloader()))\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(train_batch[0].data[i][0].numpy(), cmap=\"gray\")\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 8, figsize=(10, 5))\n",
+    "for i in range(8):\n",
+    "    ax[i].imshow(val_batch[0].data[i][0].numpy(), cmap=\"gray\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "czi",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

careamics/dataset_ng/{demo_dataset.ipynb → demos/demo_dataset.ipynb}

@@ -14,8 +14,12 @@
     "import skimage\n",
     "import tifffile\n",
     "\n",
-    "from careamics.config import create_n2n_configuration\n",
-    "from careamics.dataset_ng.dataset.dataset import CareamicsDataset, Mode"
+    "from careamics.config.configuration_factories import (\n",
+    "    _create_ng_data_configuration,\n",
+    "    _list_spatial_augmentations,\n",
+    ")\n",
+    "from careamics.dataset_ng.dataset import Mode\n",
+    "from careamics.dataset_ng.factory import create_dataset"
    ]
   },
   {
@@ -57,37 +61,36 @@
    "source": [
     "# 1. Train val from an array\n",
     "\n",
-    "train_data_config = create_n2n_configuration(\n",
-    "    \"test_exp\",\n",
+    "train_data_config = _create_ng_data_configuration(\n",
     "    data_type=\"array\",\n",
     "    axes=\"YX\",\n",
     "    patch_size=(32, 32),\n",
     "    batch_size=1,\n",
-    "    num_epochs=1,\n",
-    ").data_config\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
     "\n",
-    "val_data_config = create_n2n_configuration(\n",
-    "    \"test_exp\",\n",
+    "val_data_config = _create_ng_data_configuration(\n",
     "    data_type=\"array\",\n",
     "    axes=\"YX\",\n",
     "    patch_size=(32, 32),\n",
     "    batch_size=1,\n",
-    "    num_epochs=1,\n",
     "    augmentations=[],\n",
-    ").data_config\n",
+    ")\n",
     "\n",
     "\n",
-    "train_dataset = CareamicsDataset(\n",
-    "    data_config=train_data_config,\n",
+    "train_dataset = create_dataset(\n",
+    "    config=train_data_config,\n",
     "    mode=Mode.TRAINING,\n",
     "    inputs=[example_data],\n",
     "    targets=[segmentation],\n",
+    "    in_memory=True,\n",
     ")\n",
-    "val_dataset = CareamicsDataset(\n",
-    "    data_config=val_data_config,\n",
+    "val_dataset = create_dataset(\n",
+    "    config=val_data_config,\n",
     "    mode=Mode.VALIDATING,\n",
     "    inputs=[example_data],\n",
     "    targets=[segmentation],\n",
+    "    in_memory=True,\n",
     ")\n",
     "\n",
     "fig, ax = plt.subplots(2, 5, figsize=(10, 5))\n",
@@ -119,32 +122,37 @@
     "tifffile.imwrite(\"example_data2.tiff\", example_data[:256, :256])\n",
     "tifffile.imwrite(\"example_target2.tiff\", segmentation[:256, :256])\n",
     "\n",
-    "train_data_config = create_n2n_configuration(\n",
-    "    \"test_exp\",\n",
+    "train_data_config = _create_ng_data_configuration(\n",
     "    data_type=\"tiff\",\n",
     "    axes=\"YX\",\n",
     "    patch_size=(32, 32),\n",
     "    batch_size=1,\n",
-    "    num_epochs=1,\n",
-    ").data_config\n",
+    "    augmentations=_list_spatial_augmentations()\n",
+    ")\n",
     "\n",
-    "val_data_config = create_n2n_configuration(\n",
-    "    \"test_exp\",\n",
+    "val_data_config = _create_ng_data_configuration(\n",
     "    data_type=\"tiff\",\n",
     "    axes=\"YX\",\n",
     "    patch_size=(32, 32),\n",
     "    batch_size=1,\n",
-    "    num_epochs=1,\n",
     "    augmentations=[],\n",
-    ").data_config\n",
+    ")\n",
     "\n",
     "data = sorted(Path(\"./\").glob(\"example_data*.tiff\"))\n",
     "targets = sorted(Path(\"./\").glob(\"example_target*.tiff\"))\n",
-    "train_dataset = CareamicsDataset(\n",
-    "    data_config=train_data_config, inputs=data, targets=targets\n",
+    "train_dataset = create_dataset(\n",
+    "    config=train_data_config,\n",
+    "    mode=Mode.TRAINING,\n",
+    "    inputs=data,\n",
+    "    targets=targets,\n",
+    "    in_memory=True,\n",
     ")\n",
-    "val_dataset = CareamicsDataset(\n",
-    "    data_config=val_data_config, inputs=data, targets=targets\n",
+    "val_dataset = create_dataset(\n",
+    "    config=val_data_config,\n",
+    "    mode=Mode.VALIDATING,\n",
+    "    inputs=data,\n",
+    "    targets=targets,\n",
+    "    in_memory=True,\n",
     ")\n",
     "\n",
     "fig, ax = plt.subplots(2, 5, figsize=(10, 5))\n",
@@ -171,26 +179,33 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from careamics.config import InferenceConfig\n",
+    "from careamics.config.data import NGDataConfig\n",
     "\n",
-    "prediction_config = InferenceConfig(\n",
+    "prediction_config = NGDataConfig(\n",
     "    data_type=\"array\",\n",
-    "    tile_size=(32, 32),\n",
-    "    tile_overlap=(16, 16),\n",
+    "    patching={\n",
+    "        \"name\": \"tiled\",\n",
+    "        \"patch_size\": (32, 32),\n",
+    "        \"overlaps\": (16, 16),\n",
+    "    },\n",
     "    axes=\"YX\",\n",
-    "    image_means=(example_data.mean(),),\n",
-    "    image_stds=(example_data.std(),),\n",
-    "    tta_transforms=False,\n",
     "    batch_size=1,\n",
+    "    image_means=[example_data.mean()],\n",
+    "    image_stds=[example_data.std()],\n",
     ")\n",
-    "prediction_dataset = CareamicsDataset(\n",
-    "    data_config=prediction_config, mode=Mode.PREDICTING, inputs=[example_data]\n",
+    "\n",
+    "prediction_dataset = create_dataset(\n",
+    "    config=prediction_config,\n",
+    "    mode=Mode.PREDICTING,\n",
+    "    inputs=[example_data],\n",
+    "    targets=None,\n",
+    "    in_memory=True,\n",
     ")\n",
     "\n",
     "fig, ax = plt.subplots(1, 5, figsize=(10, 5))\n",
     "ax[0].set_title(\"Prediction input\")\n",
     "for i in range(5):\n",
-    "    sample, _ = prediction_dataset[i]\n",
+    "    sample, *_ = prediction_dataset[i]\n",
     "    ax[i].imshow(sample.data[0])"
    ]
   },
@@ -209,26 +224,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "train_data_config = create_n2n_configuration(\n",
-    "    \"test_exp\",\n",
+    "train_data_config = _create_ng_data_configuration(\n",
     "    data_type=\"custom\",\n",
     "    axes=\"YX\",\n",
     "    patch_size=(32, 32),\n",
     "    batch_size=1,\n",
-    "    num_epochs=1,\n",
-    ").data_config\n",
+    "    augmentations=_list_spatial_augmentations(),\n",
+    ")\n",
     "\n",
     "\n",
-    "def read_data_func_test(data):\n",
+    "def read_data_func_test(example_data):\n",
     "    return 255 - example_data\n",
     "\n",
     "\n",
     "fig, ax = plt.subplots(1, 5, figsize=(10, 5))\n",
-    "train_dataset = CareamicsDataset(\n",
-    "    data_config=train_data_config,\n",
+    "train_dataset = create_dataset(\n",
+    "    config=train_data_config,\n",
     "    mode=Mode.TRAINING,\n",
     "    inputs=[example_data],\n",
     "    targets=[segmentation],\n",
+    "    in_memory=True,\n",
     "    read_func=read_data_func_test,\n",
     "    read_kwargs={}\n",
     ")\n",
@@ -237,19 +252,11 @@
     "    sample, _ = train_dataset[i]\n",
     "    ax[i].imshow(sample.data[0])"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "10",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "czi",
    "language": "python",
    "name": "python3"
   },
@@ -263,7 +270,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.11"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,

careamics/dataset_ng/{demo_patch_extractor.py → demos/demo_patch_extractor.py}

@@ -1,10 +1,12 @@
 # %%
 import numpy as np
 
-# %%
-from careamics.config.support import SupportedData
-from careamics.dataset_ng.patch_extractor import create_patch_extractor
 from careamics.dataset_ng.patch_extractor.image_stack import InMemoryImageStack
+
+# %%
+from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
+    create_array_extractor,
+)
 from careamics.dataset_ng.patching_strategies import RandomPatchingStrategy
 
 # %%
@@ -30,12 +32,8 @@ print(target2)
 
 # %%
 # define example readers
-input_patch_extractor = create_patch_extractor(
-    [array1, array2], axes="SYX", data_type=SupportedData.ARRAY
-)
-target_patch_extractor = create_patch_extractor(
-    [target1, target2], axes="SYX", data_type=SupportedData.ARRAY
-)
+input_patch_extractor = create_array_extractor([array1, array2], axes="SYX")
+target_patch_extractor = create_array_extractor([target1, target2], axes="SYX")
 
 # %%
 # generate random patch specification