Complete cell-SAM with two channel selection

.gitignore

@@ -107,6 +107,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+deepcell-env/
 
 # Spyder project settings
 .spyderproject

backend/.env.example

@@ -17,3 +17,8 @@ FLASK_ENV=development
 
 # Flask monitoring dashboard
 DASHBOARD_CONFIG=fmd_config.cfg
+
+# Cellsam server
+CELLSAM_IP=127.0.0.1
+CELLSAM_PORT=8765
+CELLSAM_SERVER_VERSION=

backend/deepcell_label/client.py

@@ -0,0 +1,29 @@
+import asyncio
+import pickle
+import zlib
+
+import websockets
+from flask import current_app
+
+from deepcell_label.config import CELLSAM_IP, CELLSAM_PORT
+
+
+# connects to cell SAM server
+async def perform_send(to_send):
+    uri = f'ws://{CELLSAM_IP}:{CELLSAM_PORT}'
+
+    async with websockets.connect(uri, ping_interval=None) as websocket:
+        data = {'img': to_send}
+        pkt = zlib.compress(pickle.dumps(data))
+        await websocket.send(pkt)
+
+        pkt_received = await websocket.recv()
+
+        mask = pickle.loads(zlib.decompress(pkt_received))
+        return mask
+
+
+def send_to_server(to_send):
+    current_app.logger.info('Sent to server to generate mask for cellSAM')
+    mask = asyncio.run(perform_send(to_send))
+    return mask

backend/deepcell_label/config.py

@@ -12,6 +12,10 @@
 AWS_SECRET_ACCESS_KEY = config('AWS_SECRET_ACCESS_KEY', default='')
 S3_BUCKET = config('S3_BUCKET', default='deepcell-label-input')
 
+# Cellsam compute server
+CELLSAM_IP = config('CELLSAM_IP', default='')
+CELLSAM_PORT = config('CELLSAM_PORT', default='')
+
 TEMPLATES_AUTO_RELOAD = config('TEMPLATES_AUTO_RELOAD', cast=bool, default=True)
 
 # SQLAlchemy settings

backend/deepcell_label/label.py

@@ -13,6 +13,8 @@
 from skimage.morphology import dilation, disk, erosion, flood, square
 from skimage.segmentation import morphological_chan_vese, watershed
 
+from deepcell_label.client import send_to_server
+
 
 class Edit(object):
     """
@@ -76,6 +78,9 @@ def load(self, labels_zip):
             self.action = edit['action']
             self.height = edit['height']
             self.width = edit['width']
+            self.d1, self.d2, self.d3, self.d4 = [
+                edit[k] for k in ['d1', 'd2', 'd3', 'd4']
+            ]
             self.args = edit.get('args', None)
             # TODO: specify write mode per cell?
             self.write_mode = edit.get('writeMode', 'overlap')
@@ -102,7 +107,8 @@ def load(self, labels_zip):
         if 'raw.dat' in zf.namelist():
             with zf.open('raw.dat') as f:
                 raw = np.frombuffer(f.read(), np.uint8)
-                self.raw = np.reshape(raw, (self.width, self.height))
+                self.rawOriginal = np.reshape(raw, (self.d1, self.d2, self.d3, self.d4))
+                self.raw = self.rawOriginal[0][0]
         elif self.action in self.raw_required:
             raise ValueError(
                 f'Include raw array in raw.json to use action {self.action}.'
@@ -418,3 +424,36 @@ def action_dilate(self, cell):
         mask = self.get_mask(cell)
         dilated = dilation(mask, square(3))
         self.add_mask(dilated, cell)
+
+    def action_select_channels(self, channels):
+        self.nuclear_channel = int(channels[0])
+        self.wholecell_channel = int(channels[1])
+
+    def action_segment_all(self, channels):
+        nuclear_channel = channels[0]
+        wholecell_channel = channels[1]
+        to_send = []
+        if self.d1 == 1:
+            to_send = self.raw.reshape(self.d3, self.d4, self.d1)
+        elif self.d1 > 1:
+            nuclear = self.rawOriginal[nuclear_channel][0]
+            wholecell = self.rawOriginal[wholecell_channel][0]
+            to_send = np.stack([nuclear, wholecell], axis=-1)
+        mask = send_to_server(to_send)
+        self.labels = mask.astype(np.int32)
+        if len(self.labels.shape) == 2:
+            self.labels = np.expand_dims(np.expand_dims(self.labels, 0), 3)
+        cells = []
+        for t in range(self.labels.shape[0]):
+            for c in range(self.labels.shape[-1]):
+                for value in np.unique(self.labels[t, :, :, c]):
+                    if value != 0:
+                        cells.append(
+                            {
+                                'cell': int(value),
+                                'value': int(value),
+                                't': int(t),
+                                'c': int(c),
+                            }
+                        )
+        self.cells = cells

backend/requirements.txt

@@ -10,9 +10,10 @@ mysqlclient~=2.1.0
 numpy
 python-decouple~=3.1
 python-dotenv~=0.19.2
-python-magic~=0.4.25
-requests~=2.29.0
+python-magic
+requests
 scikit-image~=0.19.0
 sqlalchemy~=1.3.24
 tifffile
 imagecodecs
+websockets

frontend/src/Project/EditControls/EditControls.js

@@ -6,6 +6,7 @@ import CellControls from './CellControls';
 import CellTypeControls from './CellTypeControls';
 import TrackingControls from './DivisionsControls';
 import SegmentControls from './SegmentControls';
+import SegmentSamControls from './SegmentSamControls';
 
 function TabPanel(props) {
   const { children, value, index, ...other } = props;
@@ -28,14 +29,16 @@ function EditControls() {
   const value = useSelector(labelMode, (state) => {
     return state.matches('editSegment')
       ? 0
-      : state.matches('editCells')
+      : state.matches('editSegmentSam')
       ? 1
-      : state.matches('editDivisions')
+      : state.matches('editCells')
       ? 2
-      : state.matches('editCellTypes')
+      : state.matches('editDivisions')
       ? 3
-      : state.matches('editSpots')
+      : state.matches('editCellTypes')
       ? 4
+      : state.matches('editSpots')
+      ? 5
       : false;
   });
 
@@ -51,15 +54,18 @@ function EditControls() {
         <SegmentControls />
       </TabPanel>
       <TabPanel value={value} index={1}>
-        <CellControls />
+        <SegmentSamControls />
       </TabPanel>
       <TabPanel value={value} index={2}>
-        <TrackingControls />
+        <CellControls />
       </TabPanel>
       <TabPanel value={value} index={3}>
-        <CellTypeControls />
+        <TrackingControls />
       </TabPanel>
       <TabPanel value={value} index={4}>
+        <CellTypeControls />
+      </TabPanel>
+      <TabPanel value={value} index={5}>
         <SpotsControls />
       </TabPanel>
     </Box>

frontend/src/Project/EditControls/EditTabs.js

@@ -13,14 +13,16 @@ function EditTabs() {
   const value = useSelector(labelMode, (state) => {
     return state.matches('editSegment')
       ? 0
-      : state.matches('editCells')
+      : state.matches('editSegmentSam')
       ? 1
-      : state.matches('editDivisions')
+      : state.matches('editCells')
       ? 2
-      : state.matches('editCellTypes')
+      : state.matches('editDivisions')
       ? 3
-      : state.matches('editSpots')
+      : state.matches('editCellTypes')
       ? 4
+      : state.matches('editSpots')
+      ? 5
       : false;
   });
   const handleChange = (event, newValue) => {
@@ -29,15 +31,18 @@ function EditTabs() {
         labelMode.send('EDIT_SEGMENT');
         break;
       case 1:
-        labelMode.send('EDIT_CELLS');
+        labelMode.send('EDIT_SEGMENT_SAM');
         break;
       case 2:
-        labelMode.send('EDIT_DIVISIONS');
+        labelMode.send('EDIT_CELLS');
         break;
       case 3:
-        labelMode.send('EDIT_CELLTYPES');
+        labelMode.send('EDIT_DIVISIONS');
         break;
       case 4:
+        labelMode.send('EDIT_CELLTYPES');
+        break;
+      case 5:
         labelMode.send('EDIT_SPOTS');
         break;
       default:
@@ -72,6 +77,7 @@ function EditTabs() {
         variant='scrollable'
       >
         <Tab sx={{ p: 0.5, minHeight: 0 }} label='Segment' />
+        <Tab sx={{ p: 0.5, minHeight: 0 }} label='Segment-CellSAM' />
         <Tab sx={{ p: 0.5, minHeight: 0 }} label='Cells' />
         <Tab sx={{ p: 0.5, minHeight: 0 }} label='Divisions' />
         <Tab sx={{ p: 0.5, minHeight: 0 }} label='Cell Types' />

frontend/src/Project/EditControls/SegmentSamControls/ActionButtons.js

@@ -0,0 +1,22 @@
+import { useSelector } from '@xstate/react';
+import { FormLabel } from '@mui/material';
+import Box from '@mui/material/Box';
+import ButtonGroup from '@mui/material/ButtonGroup';
+import SegmentAllButton from './ActionButtons/SegmentAllButton';
+import { useRaw } from '../../ProjectContext';
+
+function ActionButtons() {
+  const raw = useRaw();
+  const layers = useSelector(raw, (state) => state.context.layers);
+  const layer = layers[0];
+  return (
+    <Box display='flex' flexDirection='column'>
+      <FormLabel>Actions</FormLabel>
+      <ButtonGroup orientation='vertical'>
+        <SegmentAllButton layer={layer} />
+      </ButtonGroup>
+    </Box>
+  );
+}
+
+export default ActionButtons;

frontend/src/Project/EditControls/SegmentSamControls/ActionButtons/ActionButton.js

@@ -0,0 +1,36 @@
+import Button from '@mui/material/Button';
+import Tooltip from '@mui/material/Tooltip';
+import { bind } from 'mousetrap';
+import React, { useEffect } from 'react';
+
+// for adding tooltip to disabled buttons
+// from https://stackoverflow.com/questions/61115913
+
+const ActionButton = ({ tooltipText, disabled, onClick, hotkey, ...other }) => {
+  const adjustedButtonProps = {
+    disabled: disabled,
+    component: disabled ? 'div' : undefined,
+    onClick: disabled ? undefined : onClick,
+  };
+
+  useEffect(() => {
+    bind(hotkey, onClick);
+  }, [hotkey, onClick]);
+
+  return (
+    <Tooltip title={tooltipText} placement='right'>
+      <Button
+        {...other}
+        {...adjustedButtonProps}
+        sx={{
+          p: 0,
+          '&.Mui-disabled': {
+            pointerEvents: 'auto',
+          },
+        }}
+      />
+    </Tooltip>
+  );
+};
+
+export default ActionButton;

frontend/src/Project/EditControls/SegmentSamControls/ActionButtons/SegmentAllButton.js

@@ -0,0 +1,92 @@
+import { useSelector } from '@xstate/react';
+import React, { useCallback } from 'react';
+import { useEditSegment, useSelect, useRaw } from '../../../ProjectContext';
+import ActionButton from './ActionButton';
+import { MenuItem, TextField } from '@mui/material';
+import Grid from '@mui/material/Grid';
+
+function LayerSelector({ layer, channelType }) {
+  const segment = useEditSegment();
+  const nuclearChannel = useSelector(segment, (state) => state.context.nuclearChannel);
+  const wholeCellChannel = useSelector(segment, (state) => state.context.wholeCellChannel);
+
+  const raw = useRaw();
+  const names = useSelector(raw, (state) => state.context.channelNames);
+
+  const onChangeNuclear = (e) => {
+    segment.send({ type: 'SET_NUCLEAR_CHANNEL', nuclearChannel: Number(e.target.value) });
+  };
+
+  const onChangeWholeCell = (e) => {
+    segment.send({ type: 'SET_WHOLE_CELL_CHANNEL', wholeCellChannel: Number(e.target.value) });
+  };
+
+  return channelType == 'nuclear' ? (
+    <TextField
+      select
+      size='small'
+      value={nuclearChannel}
+      onChange={onChangeNuclear}
+      sx={{ width: 130 }}
+    >
+      {names.map((opt, index) => (
+        <MenuItem key={index} value={index}>
+          {opt}
+        </MenuItem>
+      ))}
+    </TextField>
+  ) : (
+    <TextField
+      select
+      size='small'
+      value={wholeCellChannel}
+      onChange={onChangeWholeCell}
+      sx={{ width: 130 }}
+    >
+      {names.map((opt, index) => (
+        <MenuItem key={index} value={index}>
+          {opt}
+        </MenuItem>
+      ))}
+    </TextField>
+  );
+}
+
+function SegmentAllButton({ props, layer }) {
+  const segment = useEditSegment();
+  const grayscale = useSelector(segment, (state) => state.matches('display.grayscale'));
+
+  const onClick = useCallback(() => segment.send('SEGMENTALL'), [segment]);
+
+  const tooltipText = (
+    <span>
+      Generate all the segmentation masks <kbd>M</kbd>
+    </span>
+  );
+
+  return (
+    <Grid>
+      Select Nuclear Channel
+      <Grid item xs={10.5}>
+        <LayerSelector layer={layer} channelType={'nuclear'} />
+      </Grid>
+      Select Whole Cell Channel
+      <Grid item xs={10.5}>
+        <LayerSelector layer={layer} channelType={'wholeCell'} />
+      </Grid>
+      <Grid>
+        <ActionButton
+          {...props}
+          // disabled={!grayscale}
+          tooltipText={grayscale ? tooltipText : 'Run cell sam on one channel'}
+          onClick={onClick}
+          hotkey='m'
+        >
+          Segment All
+        </ActionButton>
+      </Grid>
+    </Grid>
+  );
+}
+
+export default SegmentAllButton;