{"cells":[{"metadata":{},"cell_type":"markdown","source":"# HPA21 Cellpose EDA\n\nI wanted to be able to try the cellpose library for cell segmentation for the HPA21 competition.\n\nUnfortunately cellpose is not one of the standard libraries avaiable to the competition notebooks\nand this presented a number of technical challenges to load the necessary library and the pre-trained\nmodel weights that the library uses without internet access (submission notebooks have to have internet\ndisabled).\n\nThis notebook demonstrates a method to offline load both the necessary libraries and the model weights.\nYou will also need to utilise the following notebooks if you want to further pursue use of cellpose.\n\n## Offline package preparation\n\n[Offline Package Wheeler](https://www.kaggle.com/andrewscholan/offline-package-wheeler-public)\n\nThis workbook is actually a general purpose notebook that imports a set of libraries using an internet\nconnection then builds the necessary wheel files so that they can be packaged into a Kaggle dataset. This\nis then used to offline load the libraries (see below).\n\n## Offline model weights preparation\n\n[Cellpose Model Collector](https://www.kaggle.com/andrewscholan/cellpose-model-collector-public)\n\nThis workbook downloads the necessary model weights used by cellpose so that they too can be\npackaged up into a Kaggle dataset."},{"metadata":{},"cell_type":"markdown","source":"# Exploratory Data Analysis\n\nThe final part of the notebook simply loads in some of the data and plots it with the masks detected by\ncellpose."},{"metadata":{},"cell_type":"markdown","source":"# Torch Custom Dataset/Rescale Transform\n\nThere is also a custom dataset for Pytorch that loads the images into suitable tensors and\na transform that is used to rescale the images before passing them to cellpose. These might prove\nuseful...?"},{"metadata":{},"cell_type":"markdown","source":"# GPU/Internet settings\n\nThis notebook should run correctly with GPU enabled and Internet turned off."},{"metadata":{"execution":{"iopub.execute_input":"2021-03-13T12:25:14.725617Z","iopub.status.busy":"2021-03-13T12:25:14.724583Z","iopub.status.idle":"2021-03-13T12:25:14.729868Z","shell.execute_reply":"2021-03-13T12:25:14.730397Z"},"papermill":{"duration":0.022057,"end_time":"2021-03-13T12:25:14.730781","exception":false,"start_time":"2021-03-13T12:25:14.708724","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"from __future__ import annotations","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.010072,"end_time":"2021-03-13T12:25:14.751793","exception":false,"start_time":"2021-03-13T12:25:14.741721","status":"completed"},"tags":[]},"cell_type":"markdown","source":"# Imports"},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2021-03-13T12:25:14.777161Z","iopub.status.busy":"2021-03-13T12:25:14.776161Z","iopub.status.idle":"2021-03-13T12:25:14.783553Z","shell.execute_reply":"2021-03-13T12:25:14.784131Z"},"papermill":{"duration":0.022038,"end_time":"2021-03-13T12:25:14.784347","exception":false,"start_time":"2021-03-13T12:25:14.762309","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import json\nimport glob\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport os\nimport PIL\nimport requests\nimport socket\n\nfrom tqdm.notebook import tqdm as show_progress\nfrom typing import Optional, Tuple, List, Any\nfrom shutil import rmtree, copytree","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import seaborn as sns","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.016728,"end_time":"2021-03-13T12:27:07.155116","exception":false,"start_time":"2021-03-13T12:27:07.138388","status":"completed"},"tags":[]},"cell_type":"markdown","source":"# Notebook configuration"},{"metadata":{"execution":{"iopub.execute_input":"2021-03-13T12:27:07.194572Z","iopub.status.busy":"2021-03-13T12:27:07.193917Z","iopub.status.idle":"2021-03-13T12:27:07.19779Z","shell.execute_reply":"2021-03-13T12:27:07.197276Z"},"papermill":{"duration":0.025898,"end_time":"2021-03-13T12:27:07.197936","exception":false,"start_time":"2021-03-13T12:27:07.172038","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"NOTEBOOK_NAME = \"HPA21 Cellpose EDA\"\n\nQUICK_TEST = True\n\nTRAIN_SLICE = slice(1000) if QUICK_TEST else slice(None)\nTEST_SLICE = slice(1000) if QUICK_TEST else slice(None)\n\n# These are dataframe column names\nID = \"ID\"\nIMAGE_HEIGHT = \"ImageHeight\"\nIMAGE_WIDTH = \"ImageWidth\"\nIMAGE_SHAPE = \"ImageShape\"\nLABEL = \"Label\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# List of the dataset channel names\nPROTEIN_OF_INTEREST = \"green\"\nNUCLEUS = \"blue\"\nMICROTUBULE = \"red\"\nENDOPLASMIC_RETICULUM = \"yellow\"\nCHANNEL_NAMES = {\n    PROTEIN_OF_INTEREST: \"Protein of interest\", \n    NUCLEUS: \"Nucleus\", \n    MICROTUBULE: \"microtubule\",\n    ENDOPLASMIC_RETICULUM: \"endoplasmic_reticulum\",\n}\n\n# List of the label names\nLABELS = [\n    \"Nucleoplasm\",\n    \"Nuclear membrane\",\n    \"Nucleoli\",\n    \"Nucleoli fibrillar center\",\n    \"Nuclear speckles\",\n    \"Nuclear bodies\",\n    \"Endoplasmic reticulum\",\n    \"Golgi apparatus\",\n    \"Intermediate filaments\",\n    \"Actin filaments\",\n    \"Microtubules\",\n    \"Mitotic spindle\",\n    \"Centrosome\",\n    \"Plasma membrane\",\n    \"Mitochondria\",\n    \"Aggresome\",\n    \"Cytosol\",\n    \"Vesicles and punctate cytosolic patterns\",\n    \"Negative\",\n]\nNUM_LABELS = len(LABELS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Make plots larger\nplt.rcParams['figure.figsize'] = [16, 9]\nplt.rcParams['figure.dpi'] = 120","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Folder Roots"},{"metadata":{"execution":{"iopub.execute_input":"2021-03-13T12:27:07.283874Z","iopub.status.busy":"2021-03-13T12:27:07.282865Z","iopub.status.idle":"2021-03-13T12:27:07.288875Z","shell.execute_reply":"2021-03-13T12:27:07.287463Z"},"papermill":{"duration":0.030209,"end_time":"2021-03-13T12:27:07.28921","exception":false,"start_time":"2021-03-13T12:27:07.259001","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"INPUT_ROOT = os.path.abspath(os.path.join(\"..\", \"input\"))\nTEMP_ROOT = os.path.abspath(os.path.join(\"..\", \"temp\"))\nWORKING_ROOT = os.path.abspath(\".\")\nprint(f\"INPUT_ROOT='{INPUT_ROOT}'\")\nprint(f\"TEMP_ROOT='{TEMP_ROOT}'\")\nprint(f\"WORKING_ROOT='{WORKING_ROOT}'\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"HPA21_DATASET_PATH = os.path.join(INPUT_ROOT, \"hpa-single-cell-image-classification\")\nprint(f\"HPA21_DATASET_PATH='{HPA21_DATASET_PATH}'\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Offline load additional packages\n\nSee [Offline Package Wheeler](https://www.kaggle.com/andrewscholan/offline-package-wheeler-public)"},{"metadata":{"trusted":true},"cell_type":"code","source":"EXTRA_PACKAGES_ROOT = os.path.join(INPUT_ROOT, \"hpa21-extra-packages\")\nrequirements_txt = os.path.join(EXTRA_PACKAGES_ROOT, \"requirements.txt\")\nwheels_path = os.path.join(EXTRA_PACKAGES_ROOT, \"wheels\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with open(requirements_txt, \"r\") as f:\n    requirements = f.readlines()\n    for requirement in requirements:\n        print(requirement.strip())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install \\\n    --requirement {requirements_txt} \\\n    --no-index \\\n    --find-links file://{wheels_path}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Offline load the Cellpose model weights\n\nSee [Cellpose Model Collector](https://www.kaggle.com/andrewscholan/cellpose-model-collector-public)\n\nThis copies the model weights to the cache folder that cellpose uses so that it doesn't\nattempt to download them from a non-existant internet connection"},{"metadata":{"trusted":true},"cell_type":"code","source":"cellpose_cache_path = os.path.join(\"/\", \"root\", \".cellpose\")\ncellpose_cache_path","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cellpose_models_as_dataset = os.path.join(INPUT_ROOT, \"cellpose-models\", \"cellpose_models\")\ncellpose_models_as_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists(cellpose_cache_path):\n    copytree(cellpose_models_as_dataset, cellpose_cache_path)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## And import the modules"},{"metadata":{"execution":{"iopub.execute_input":"2021-03-13T12:27:06.222496Z","iopub.status.busy":"2021-03-13T12:27:06.221707Z","iopub.status.idle":"2021-03-13T12:27:07.077693Z","shell.execute_reply":"2021-03-13T12:27:07.07867Z"},"papermill":{"duration":0.881083,"end_time":"2021-03-13T12:27:07.078903","exception":false,"start_time":"2021-03-13T12:27:06.19782","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport torchaudio\nimport cellpose\nprint(f\"pytorch=={torch.__version__}\")\nprint(f\"torchvision=={torchvision.__version__}\")\nprint(f\"torchaudio=={torchaudio.__version__}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Set torchaudio back-end for cross-platform use\ntorchaudio.USE_SOUNDFILE_LEGACY_INTERFACE = False\ntorchaudio.set_audio_backend(\"soundfile\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from cellpose import io as cellpose_io\nimport torchvision.transforms.functional as TF","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.017781,"end_time":"2021-03-13T12:27:07.756249","exception":false,"start_time":"2021-03-13T12:27:07.738468","status":"completed"},"tags":[]},"cell_type":"markdown","source":"# GPU detection/setup"},{"metadata":{"execution":{"iopub.execute_input":"2021-03-13T12:27:07.807991Z","iopub.status.busy":"2021-03-13T12:27:07.798304Z","iopub.status.idle":"2021-03-13T12:27:08.534484Z","shell.execute_reply":"2021-03-13T12:27:08.533804Z"},"papermill":{"duration":0.760353,"end_time":"2021-03-13T12:27:08.534676","exception":false,"start_time":"2021-03-13T12:27:07.774323","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# Execute the following command. If we are in a notebook on a system with\n# an NVIDIA GPU it will give us the CUDA Version. This needs to match the\n# version of PyTorch that we installed above.\n# If we are not on a GPU enabled host then it will say command not found.\n!nvidia-smi","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2021-03-13T12:27:08.58615Z","iopub.status.busy":"2021-03-13T12:27:08.585436Z","iopub.status.idle":"2021-03-13T12:27:08.589461Z","shell.execute_reply":"2021-03-13T12:27:08.589961Z"},"papermill":{"duration":0.03672,"end_time":"2021-03-13T12:27:08.590171","exception":false,"start_time":"2021-03-13T12:27:08.553451","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import torch.cuda\nuse_gpu = torch.cuda.is_available()\ndevice_math = torch.device(type=\"cuda\" if use_gpu else \"cpu\")\ndevice_cpu = torch.device(\"cpu\")\nif device_math.type==\"cuda\":\n    print(f\"Using GPU when tensors are loaded to 'device_math'.\\n\"\n          f\"  Device: {torch.cuda.get_device_name()}\\n\"\n          f\"  Number of GPUs: {torch.cuda.device_count()}\\n\"\n          f\"  GPU initialised: {torch.cuda.is_initialized()}\\n\")\n    print(torch.cuda.memory_summary())\nelse:\n    print(\"GPU not available, all tensors loaded to 'device_math' will reside \"\n          \"in CPU memory.\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualise some images"},{"metadata":{},"cell_type":"markdown","source":"## Load in the CSV file"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load the training CSV\ntrain_csv = os.path.join(HPA21_DATASET_PATH, \"train.csv\")\ntrain_df = pd.read_csv(train_csv)\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## One-hot encode the labels"},{"metadata":{"trusted":true},"cell_type":"code","source":"labels = train_df[LABEL].str.split(\"|\")\nlabels.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"has_label = lambda label_list, label_num: 1.0 if str(label_num) in label_list else 0.0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_truth = {label_num: labels.apply(has_label, label_num=label_num)\n               for label_num in range(NUM_LABELS)}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_truth[ID] = train_df[ID]\nlabel_truth[ID].head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_truth = pd.DataFrame(label_truth).set_index(ID)\ntrain_truth.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_truth_melted = train_truth.melt(value_vars=train_truth.columns)\ntrain_truth_melted = train_truth_melted[train_truth_melted[\"value\"] != 0.0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sns.catplot(x=\"variable\", kind=\"count\", data=train_truth_melted, \n            order=train_truth_melted[\"variable\"].value_counts().index);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Get the image height and width"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Function that computes the image height and width for all images\n# in a folder.\ndef read_image_dimensions(\n    folder_path: str, \n    file_ext: str = \".png\",\n    id: str = ID, \n    height: str = IMAGE_HEIGHT, \n    width: str = IMAGE_WIDTH,\n    shape: str = IMAGE_SHAPE,\n    channel: str = PROTEIN_OF_INTEREST,\n    file_slice: slice = slice(None),\n) -> pd.DataFrame:\n    \"\"\"\n    Reads the image height and width for all files in the folder path\n    that end in with the string specified in channel.\n    \n    Args:\n        folder_path: The path to the folder containing the image files.\n        file_ext: The file extension that we are interested in.\n        id: The column name for the file id. The file id is acquired from\n            the filename which is assumed to be in the format \n            <id>_<channel>.ext.\n        height: The column name for the image height.\n        width: The column name for the image width.\n        shape: The column name for the image shape.\n        channel: The channel image to use for sizing (see id).\n        slice: Slice to apply to the file list (used for testing).\n        \n    Returns:\n        Dataframe consisting of columns {id}, {width}. {height} and {[channel_path]} indexed\n        by {id}.\n    \"\"\"\n    # Get all the file paths that match the channel\n    file_filter = os.path.join(folder_path, f\"*_{channel}{file_ext}\")\n    file_paths = glob.glob(file_filter)[file_slice]\n    \n    image_size: Tuple[int, int] = lambda file_path: PIL.Image.open(file_path).size\n    def file_id(file_path: str) -> str:\n        _, file_name = os.path.split(file_path)\n        file_name, _ = os.path.splitext(file_name)\n        id, _ = file_name.split(\"_\")\n        return id\n\n    ids = [file_id(file_path) for file_path in file_paths]\n    widths_and_heights = [image_size(file_path) for file_path in show_progress(file_paths)]\n    \n    df_dict = {\n        id: ids,\n        width: [image_width for image_width, _ in widths_and_heights],\n        height: [image_height for _, image_height in widths_and_heights],\n        shape: [f\"W{image_width}xH{image_height}\" \n                for image_width, image_height in widths_and_heights]\n    }\n    \n    return pd.DataFrame(df_dict).set_index(id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_path = os.path.join(HPA21_DATASET_PATH, \"train\")\ntrain_image_sizes = read_image_dimensions(train_path, file_slice=TRAIN_SLICE)\ntrain_image_sizes.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_path = os.path.join(HPA21_DATASET_PATH, \"test\")\ntest_image_sizes = read_image_dimensions(test_path, file_slice=TEST_SLICE)\ntest_image_sizes.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_image_shapes = set(train_image_sizes[IMAGE_SHAPE].unique())\ntest_image_shapes = set(test_image_sizes[IMAGE_SHAPE].unique())\nimage_shapes = sorted(list(train_image_shapes | test_image_shapes))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sns.catplot(x=IMAGE_SHAPE, kind=\"count\", data=train_image_sizes,\n            order=image_shapes);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sns.catplot(x=IMAGE_SHAPE, kind=\"count\", data=test_image_sizes,\n            order=image_shapes);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Add shapes to truth dataframes"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get the subset of the truth that corresponds with the slice that we are using\n# for this note-book by using the index of the images we sized\ntrain_truth_slice = train_truth.loc[train_image_sizes.index, :]\ntrain = pd.concat([train_truth_slice, train_image_sizes], axis=\"columns\")\ntrain","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# And given that the test data is not labelled then our \"truth\" for the test\n# data is the image size dataframe\ntest = test_image_sizes\ntest","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Transform to rescale the images\nThe images are very high resolution, but I suspect that for cell segmentation we don't need to use them that size."},{"metadata":{"trusted":true},"cell_type":"code","source":"class RescaleImage(object):\n    \"\"\"\n    Transform callable class that can be used to transform the tensor size by a\n    scale value. Under the hood it uses the torchvision.functional.resize transform\n    to do the work.\n    \"\"\"\n    \n    def __init__(self, scale:int = 0.25, interpolation:int = PIL.Image.BILINEAR):\n        \"\"\"\n        This initailises the transform function.\n        \n        Args:\n            scale: The scale factor for the image W and H dimensions (E.g. 0.5 will\n                reduce the image dimensions by 50% in W and H)\n            interpolation: The interpolation technique to use for the rescaling.\n        \"\"\"\n        self._scale = scale\n        self._interpolation = interpolation\n        \n    def __call__(self, t:torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        This is the transforming function. It takes as input a tensor or shape\n        [..., H, W] and returns a tensor of the same shape but with dimensions H and\n        W scaled as set in the initialiser.\n        \n        Args:\n            t: Input tensor of shape [..., H, W]\n            \n        Returns:\n            Tensor of shape [..., H, W] with scaled dimensions\n        \"\"\"\n        shape_in = t.size()\n        height, width = list(shape_in)[-2:]\n        new_size = [int(height * self._scale), int(width * self._scale)]\n        t_out = TF.resize(\n            img=t, size=new_size, interpolation=self._interpolation\n        )\n        return t_out\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data loader\n\nThis is a data loader that takes images of the same size and combines the individual images into composites."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Dataset for HPA21 images.\nclass CellImageDataset(torch.utils.data.Dataset):\n    \"\"\"\n    Defines an image dataset for the HPA21 images.\n    \"\"\"\n    \n    def __init__(\n        self, \n        images: pd.DataFrame, \n        root_dir: str,\n        file_ext: str = \".png\",\n        id: str = ID,\n        channels: List[List[str]] = [\n            [ENDOPLASMIC_RETICULUM, NUCLEUS, MICROTUBULE], \n            [PROTEIN_OF_INTEREST]\n        ],\n        target_cols=None, \n        transform=None, \n        transform_target=None\n    ) -> List[torch.Tensor]:\n        \"\"\"\n        Creates the dataset based on the images dataframe.\n        \n        Args:\n            images: Dataframe indexed by id; the id is the root of the filename.\n            root_dir: The directory containing all of the images.\n            id: The name of the ID column.\n            channels: List of tesnsors to return and the channels to include in each tensor\n            target_cols: List of column names that constitututes the target labels. Set\n                to None if there are no labels.\n            transform: Transform to perform on the image after loading.\n            transform_target: Transform to perform on the target tensor.\n        \"\"\"\n        super().__init__()\n        self._images = images.reset_index()\n        self._root_dir = root_dir\n        self._file_ext = file_ext\n        self._id = id\n        self._channels = channels\n        self._target_cols = target_cols\n        self._transform = transform\n        self._transform_target = transform_target\n        \n    def __len__(self):\n        \"\"\"\n        Return the size of the dataset. This is defined by the dataframe that was input.\n        \n        Returns:\n            The size of the dataset (based on the dataframe passed in).\n        \"\"\"\n        return len(self._images)\n    \n    def __getitem__(self, idx):\n        \"\"\"\n        Gets the data associated with a specific index.\n        \n        Args:\n            idx: Index of data to get.\n            \n        Returns:\n            A list of tensors; the data tensors are defined by the channels\n            parameter; the labels tensor is appended if target_cols is specified.\n        \"\"\"\n        # Get this image ID\n        this_image = self._images.iloc[idx, :]\n        this_image_id = this_image[self._id]\n        \n        # Build the filename lists to load and combine\n        channel_file_paths = [\n            [\n                os.path.join(\n                    self._root_dir, f\"{this_image_id}_{channel}{self._file_ext}\"\n                )\n                for channel in tensor_channels\n            ]\n            for tensor_channels in self._channels\n        ]\n        \n        # Now load the images as numpy arrays using the cellpose_io imread\n        channel_image_arrays = [\n            np.stack(\n                [\n                    cellpose_io.imread(file_path)       \n                    for file_path in tensor_channel_file_paths\n                ]\n            )\n            for tensor_channel_file_paths in channel_file_paths\n        ]\n        \n        # Convert to torch tensors\n        channel_tensors = [\n            torch.tensor(channel_image_array.astype(np.float32) / 256.0)\n            for channel_image_array in channel_image_arrays\n        ]\n        \n        # Now apply any transforms to them\n        if self._transform is not None:\n            channel_tensors = [self._transform(channel_tensor) \n                               for channel_tensor in channel_tensors]\n        \n        # Now process the targets, if necessary\n        if self._target_cols is not None:\n            # Extract the labels as a numpy array\n            targets = this_image.loc[target_cols].values.astype(np.float32)\n            # Convert to a tensor and shape to be CT\n            target_tensors = [torch.tensor(targets).reshape(1, -1)]\n        else:\n            target_tensors = []\n    \n        # And apply the target transform\n        if self._transform_target is not None:\n            target_tensors = [self._transform_target(target_tensor)\n                              for target_tensor in target_tensors]\n    \n        return tuple(channel_tensors + target_tensors)\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"target_cols = list(range(NUM_LABELS))\ntrain_dataset = CellImageDataset(\n    images = train,\n    root_dir = train_path,\n    target_cols = target_cols,\n    transform = RescaleImage()\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build the datasets by image label - for visualisations"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_datasets = {\n    label: CellImageDataset(\n        images = train[train[label]==1.0],\n        root_dir = train_path,\n        target_cols = target_cols,\n        transform = RescaleImage()\n    )\n    for label in target_cols\n}\ntrain_datasets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for label, dataset in train_datasets.items():\n    print(f\"{label}: {len(dataset)}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Plot an image for each label type"},{"metadata":{"trusted":true},"cell_type":"code","source":"cell_image, protein_image, labels = next(iter(train_datasets[0]))\ncell_image.size()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_images = {label: next(iter(train_datasets[label]))\n               for label in show_progress(target_cols)\n               if len(train_datasets[label])>0}\n\nmake_im = lambda t: t.detach().cpu().numpy().transpose((1, 2, 0))\n\ncell_images = {label: make_im(image)\n               for label, (image, _, _) in plot_images.items()}\nprotein_images = {label: make_im(image)\n                  for label, (_, image, _) in plot_images.items()}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for index, label in enumerate(sorted(plot_images.keys())):\n    label_name = LABELS[label]\n    cell_image = cell_images[label]\n    protein_image = protein_images[label]\n    \n    plt.subplot(1, 4, 1)\n    plt.imshow(cell_image, origin=\"lower\")\n    \n    plt.subplot(1, 4, 3)\n    plt.imshow(protein_image, cmap=\"Greens\", origin=\"lower\")\n    plt.title(f\"{label}: {label_name}\")\n    \n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Segmentation with Cellpose"},{"metadata":{"trusted":true},"cell_type":"code","source":"from cellpose import models\n\n# DEFINE CELLPOSE MODEL\n# model_type='cyto' or model_type='nuclei'\n# We want the cell extent so set it as cyto\nmodel = models.Cellpose(gpu=use_gpu, model_type='cyto', torch=True)\n\n# define CHANNELS to run segementation on\n# grayscale=0, R=1, G=2, B=3\n# channels = [cytoplasm, nucleus]\n# if NUCLEUS channel does not exist, set the second channel to 0\n# channels = [0,0]\n# IF ALL YOUR IMAGES ARE THE SAME TYPE, you can give a list with 2 elements\n# channels = [0,0] # IF YOU HAVE GRAYSCALE\n# channels = [2,3] # IF YOU HAVE G=cytoplasm and B=nucleus\n# channels = [2,1] # IF YOU HAVE G=cytoplasm and R=nucleus\nchannels = [1, 2]\n\n# if diameter is set to None, the size of the cells is estimated on a per image basis\n# you can set the average cell `diameter` in pixels yourself (recommended) \n# diameter can be a list or a single number for all images\ndiameter = None\n\nimages = list(cell_images.values())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# you can run all in a list e.g.\nmasks, flows, styles, diams = model.eval(images, diameter=diameter, channels=channels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for index, label in enumerate(sorted(plot_images.keys())):\n    label_name = LABELS[label]\n    cell_image = cell_images[label]\n    protein_image = protein_images[label]\n    mask = masks[index]\n    \n    plt.subplot(1, 4, 1)\n    plt.imshow(cell_image, origin=\"lower\")\n    \n    plt.subplot(1, 4, 2)\n    mask_image = cellpose.plot.mask_rgb(mask)\n    plt.imshow(mask_image, origin=\"lower\")\n    \n    plt.subplot(1, 4, 3)\n    plt.title(f\"{label}: {label_name}\")\n    plt.imshow(protein_image, cmap=\"Greens\", origin=\"lower\")\n    \n    plt.subplot(1, 4, 4)\n    mask_image = cellpose.plot.mask_overlay(protein_image, mask)\n    plt.imshow(mask_image, origin=\"lower\")\n    \n    plt.show()\n","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}