Implement cell types machine, basic controls, and canvas rendering (#392

)

* Add cellTypesMachine.js and event bus

* Add cell types controls prototype

* Add export and canvas

Cell types now exported and can be visualized on a dedicated canvas by color

* Implement adding cell mode

One can enter and escape the addingCells mode and add multiple cells before exiting

* Fix loadMachine to default to empty cell types list

Fixes a bug where old project zips, for example, would crash because no cellTypes.json is inside the zip file loaded on the front end

* Apply automatic changes

---------

Co-authored-by: ykevu <ykevu@users.noreply.github.com>

backend/deepcell_label/loaders.py

@@ -35,6 +35,7 @@ def __init__(self, image_file=None, label_file=None, axes=None):
         self.y = None
         self.spots = None
         self.divisions = None
+        self.cellTypes = None
         self.cells = None
 
         self.image_file = image_file
@@ -55,6 +56,7 @@ def load(self):
         self.y = load_segmentation(self.label_file)
         self.spots = load_spots(self.label_file)
         self.divisions = load_divisions(self.label_file)
+        self.cellTypes = load_cellTypes(self.label_file)
         self.cells = load_cells(self.label_file)
 
         if self.y is None:
@@ -67,6 +69,7 @@ def write(self):
         self.write_segmentation()
         self.write_spots()
         self.write_divisions()
+        self.write_cellTypes()
         self.write_cells()
 
     def write_images(self):
@@ -119,6 +122,10 @@ def write_divisions(self):
         """Writes divisions to divisions.json in the output zip."""
         self.zip.writestr('divisions.json', json.dumps(self.divisions))
 
+    def write_cellTypes(self):
+        """Writes cell types to cellTypes.json in the output zip."""
+        self.zip.writestr('cellTypes.json', json.dumps(self.cellTypes))
+
     def write_cells(self):
         """Writes cells to cells.json in the output zip."""
         if self.cells is None:
@@ -227,6 +234,26 @@ def load_divisions(f):
     return divisions
 
 
+def load_cellTypes(f):
+    """
+    Load cell types from cellTypes.json in project archive
+
+    Args:
+        zf: zip file with cellTypes.json
+
+    Returns:
+        dict or None if cellTypes.json not found
+    """
+    f.seek(0)
+    cellTypes = None
+    if zipfile.is_zipfile(f):
+        zf = zipfile.ZipFile(f, 'r')
+        cellTypes = load_zip_json(zf, filename='cellTypes.json')
+    if cellTypes is None:
+        return []
+    return cellTypes
+
+
 def load_cells(f):
     """
     Load cells from label file.

frontend/package.json

@@ -49,6 +49,7 @@
     "quickselect": "^2.0.0",
     "react": "^17.0.1",
     "react-archer": "^3.3.0",
+    "react-color": "^2.19.3",
     "react-dom": "^17.0.1",
     "react-dropzone": "^11.4.0",
     "react-error-boundary": "^3.1.3",

frontend/src/Project/Canvas/ToolCanvas/AddCellTypeCanvas.js

@@ -0,0 +1,20 @@
+// Canvas when adding cells to a cell type
+
+import { useSelector } from '@xstate/react';
+import { useEditCellTypes } from '../../ProjectContext';
+import OutlineCellCanvas from './OutlineCellCanvas';
+
+const white = [1, 1, 1, 1];
+
+function AddCellTypeCanvas({ setBitmaps }) {
+  const editCellTypes = useEditCellTypes();
+  const cell = useSelector(editCellTypes, (state) => state.context.cell);
+
+  if (!cell) {
+    return null;
+  }
+
+  return <OutlineCellCanvas setBitmaps={setBitmaps} cell={cell} color={white} />;
+}
+
+export default AddCellTypeCanvas;

frontend/src/Project/Canvas/ToolCanvas/CellTypeCanvas.js

@@ -0,0 +1,125 @@
+// Canvas when visualizing cell types
+import { useSelector } from '@xstate/react';
+import { useEffect, useRef } from 'react';
+import {
+  useAlphaGpu,
+  useArrays,
+  useCanvas,
+  useCellMatrix,
+  useCellTypes,
+  useImage,
+  useLabeled,
+  useSelectedCell,
+} from '../../ProjectContext';
+
+function CellTypeCanvas({ setBitmaps }) {
+  const canvas = useCanvas();
+  const width = useSelector(canvas, (state) => state.context.width);
+  const height = useSelector(canvas, (state) => state.context.height);
+
+  const labeled = useLabeled();
+  const feature = useSelector(labeled, (state) => state.context.feature);
+
+  const image = useImage();
+  const t = useSelector(image, (state) => state.context.t);
+
+  const arrays = useArrays();
+  const labeledArray = useSelector(
+    arrays,
+    (state) => state.context.labeled && state.context.labeled[feature][t]
+  );
+
+  const cell = useSelectedCell();
+
+  const cellTypes = useCellTypes();
+  const colorMap = useSelector(cellTypes, (state) => state.context.colorMap);
+
+  const cellMatrix = useCellMatrix();
+
+  const gpu = useAlphaGpu();
+  const kernelRef = useRef();
+
+  useEffect(() => {
+    const kernel = gpu.createKernel(
+      `function (data, cell, cells, numLabels, numValues, colorMap) {
+        const x = this.thread.x;
+        const y = this.constants.h - 1 - this.thread.y;
+        const value = data[y][x];
+        let north = value;
+        let south = value;
+        let east = value;
+        let west = value;
+        if (x !== 0) {
+          north = data[y][x - 1];
+        }
+        if (x !== this.constants.w - 1) {
+          south = data[y][x + 1];
+        }
+        if (y !== 0) {
+          west = data[y - 1][x];
+        }
+        if (y !== this.constants.h - 1) {
+          east = data[y + 1][x];
+        }
+        let outlineOpacity = 1;
+        if (value < numValues) {
+          for (let i = 0; i < numLabels; i++) {
+            if (cells[value][i] === 1) {
+              if (cells[north][i] === 0 || cells[south][i] === 0 || cells[west][i] === 0 || cells[east][i] === 0
+                  || north >= numValues || south >= numValues || west >= numValues || east >= numValues)
+              {
+                if (i === cell) {
+                  this.color(1, 1, 1, 1);
+                }
+                else {
+                  const [r, g, b, a] = colorMap[i];
+                  this.color(r, g, b, a);
+                }
+              }
+              else {
+                if (colorMap[i][3] !== 0) {
+                  const [r, g, b, a] = colorMap[i];
+                  this.color(r, g, b, 0.3);
+                }
+              }
+            }
+          }
+        }
+      }`,
+      {
+        constants: { w: width, h: height },
+        output: [width, height],
+        graphical: true,
+        dynamicArguments: true,
+      }
+    );
+    kernelRef.current = kernel;
+  }, [gpu, width, height]);
+
+  useEffect(() => {
+    const kernel = kernelRef.current;
+
+    if (labeledArray && cellMatrix) {
+      const numValues = cellMatrix.length;
+      const numLabels = cellMatrix[0].length;
+      kernel(labeledArray, cell, cellMatrix, numLabels, numValues, colorMap);
+      // Rerender the parent canvas
+      createImageBitmap(kernel.canvas).then((bitmap) => {
+        setBitmaps((bitmaps) => ({ ...bitmaps, types: bitmap }));
+      });
+    }
+  }, [labeledArray, cell, cellMatrix, colorMap, setBitmaps, width, height]);
+
+  useEffect(
+    () => () =>
+      setBitmaps((bitmaps) => {
+        const { types, ...rest } = bitmaps;
+        return rest;
+      }),
+    [setBitmaps]
+  );
+
+  return null;
+}
+
+export default CellTypeCanvas;

frontend/src/Project/Canvas/ToolCanvas/ToolCanvas.js

@@ -1,8 +1,16 @@
 import { useSelector } from '@xstate/react';
 import React from 'react';
-import { useEditCells, useEditDivisions, useEditSegment, useLabelMode } from '../../ProjectContext';
+import {
+  useEditCells,
+  useEditCellTypes,
+  useEditDivisions,
+  useEditSegment,
+  useLabelMode,
+} from '../../ProjectContext';
 import AddDaughterCanvas from './AddDaughterCanvas';
+import AddCellTypeCanvas from './AddCellTypeCanvas';
 import BrushCanvas from './BrushCanvas';
+import CellTypeCanvas from './CellTypeCanvas';
 import FloodCanvas from './FloodCanvas';
 import ReplaceCanvas from './ReplaceCanvas';
 import SwapCanvas from './SwapCanvas';
@@ -19,6 +27,9 @@ function ToolCanvas({ setBitmaps }) {
   const editDivisions = useEditDivisions();
   const addingDaughter = useSelector(editDivisions, (state) => state.matches('addingDaughter'));
 
+  const editCellTypes = useEditCellTypes();
+  const addingCell = useSelector(editCellTypes, (state) => state.matches('addingCell'));
+
   const labelMode = useLabelMode();
   const mode = useSelector(labelMode, (state) =>
     state.matches('editSegment')
@@ -27,6 +38,8 @@ function ToolCanvas({ setBitmaps }) {
       ? 'cells'
       : state.matches('editDivisions')
       ? 'divisions'
+      : state.matches('editCellTypes')
+      ? 'cellTypes'
       : false
   );
 
@@ -58,6 +71,16 @@ function ToolCanvas({ setBitmaps }) {
         return <AddDaughterCanvas setBitmaps={setBitmaps} />;
       }
       return null;
+    case 'cellTypes':
+      if (addingCell) {
+        return (
+          <>
+            <CellTypeCanvas setBitmaps={setBitmaps} />
+            <AddCellTypeCanvas setBitmaps={setBitmaps} />
+          </>
+        );
+      }
+      return <CellTypeCanvas setBitmaps={setBitmaps} />;
     default:
       return null;
   }

frontend/src/Project/DisplayControls/Cells/Selected/Selected.js

@@ -1,7 +1,7 @@
 import Box from '@mui/material/Box';
 import { useSelector } from '@xstate/react';
 import React, { useEffect, useState } from 'react';
-import { useSelect } from '../../../ProjectContext';
+import { useEditCellTypes, useSelect } from '../../../ProjectContext';
 import Cell from '../Cell';
 import NewCellButton from './NewCellButton';
 import NextCellButton from './NextCellButton';
@@ -12,13 +12,15 @@ let numMounted = 0;
 
 function Selected() {
   const select = useSelect();
+  const editCellTypes = useEditCellTypes();
   const cell = useSelector(select, (state) => state.context.selected);
 
   useEffect(() => {
     const listener = (e) => {
       switch (e.key) {
         case 'Escape':
           select.send('RESET');
+          editCellTypes.send('RESET');
           break;
         case 'n':
           select.send('SELECT_NEW');

frontend/src/Project/EditControls/CellTypeControls/CellTypeControls.js

@@ -0,0 +1,33 @@
+// Controls for cell types menu, including button for adding cell type,
+// the list of cell types, and an editing prompt when adding cells
+
+import { Box, FormLabel, Button } from '@mui/material';
+import AddCellTypeLabel from './CellTypeUI/AddCellTypeLabel';
+import CellTypeAccordionList from './CellTypeUI/CellTypeAccordionList';
+import EditingPrompt from './CellTypeUI/EditingPrompt';
+
+function CellTypeControls() {
+  return (
+    <Box display='flex' flexDirection='column'>
+      <FormLabel sx={{ marginBottom: 2 }}>
+        Cell Type Labels
+        <Button
+          variant='contained'
+          disableElevation
+          disableRipple
+          style={{ borderRadius: 100 }}
+          color='secondary'
+          sx={{ width: 5, height: 20, top: -1, marginLeft: 1 }}
+        >
+          {' '}
+          Beta
+        </Button>
+      </FormLabel>
+      <AddCellTypeLabel />
+      <CellTypeAccordionList />
+      <EditingPrompt />
+    </Box>
+  );
+}
+
+export default CellTypeControls;

frontend/src/Project/EditControls/CellTypeControls/CellTypeUI/AddCellTypeLabel.js

@@ -0,0 +1,60 @@
+import AddIcon from '@mui/icons-material/Add';
+import { Box, Button } from '@mui/material';
+import ClickAwayListener from '@mui/material/ClickAwayListener';
+import Popper from '@mui/material/Popper';
+import { useReducer, useRef } from 'react';
+import { TwitterPicker } from 'react-color';
+import { useEditCellTypes } from '../../../ProjectContext';
+
+function AddCellTypeLabel() {
+  const editCellTypesRef = useEditCellTypes();
+
+  const [open, toggle] = useReducer((v) => !v, false);
+  const anchorRef = useRef(null);
+
+  const handleChange = (color) => {
+    editCellTypesRef.send({ type: 'ADD_TYPE', color: color.hex });
+    toggle();
+  };
+
+  return (
+    <Box sx={{ position: 'relative', boxShadow: 3 }}>
+      <Button
+        variant='contained'
+        startIcon={
+          <AddIcon
+            sx={{
+              position: 'relative',
+              top: -0.3,
+              left: -5,
+            }}
+          />
+        }
+        style={{ borderRadius: 10 }}
+        sx={{
+          position: 'absolute',
+          m: 1,
+          my: 0.25,
+          top: 0,
+          margin: 'auto',
+          height: '2.5rem',
+          width: 300,
+          fontWeight: 'bold',
+          fontSize: 14,
+        }}
+        onClick={toggle}
+        ref={anchorRef}
+        size='large'
+      >
+        <div style={{ marginRight: 106 }}>Add Cell Type</div>
+      </Button>
+      <Popper open={open} anchorEl={anchorRef.current} placement='bottom'>
+        <ClickAwayListener onClickAway={toggle}>
+          <TwitterPicker onChange={handleChange} />
+        </ClickAwayListener>
+      </Popper>
+    </Box>
+  );
+}
+
+export default AddCellTypeLabel;