Changed channel selection method to add secondary channel to primary channel

.gitignore

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

backend/deepcell_label/client.py

@@ -0,0 +1,28 @@
+import asyncio
+import os
+import pickle
+import zlib
+
+import websockets
+from flask import current_app
+
+
+# connects to cell SAM server
+async def perform_send(to_send):
+    uri = os.environ['CELLSAM_SERVER']  # to be replaced with cell SAM server uri
+    async with websockets.connect(uri, ping_interval=None) as websocket:
+        data = {'img': to_send}
+        print(uri)
+        pkt = zlib.compress(pickle.dumps(data))
+        await websocket.send(pkt)
+        print('sent')
+        pkt_received = await websocket.recv()
+        print('received')
+        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/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,32 @@ def action_dilate(self, cell):
         mask = self.get_mask(cell)
         dilated = dilation(mask, square(3))
         self.add_mask(dilated, cell)
+
+    def action_segment_all(self, channels):
+        primary_channel = channels[0]
+        secondary_channel = channels[1]
+        to_send = []
+        if self.d1 == 1:
+            to_send = self.raw.reshape(self.d3, self.d4, self.d1)
+        elif self.d1 > 1:
+            primary = self.rawOriginal[primary_channel][0]
+            secondary = self.rawOriginal[secondary_channel][0]
+            to_send = np.stack([primary, secondary], 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
+# python-magic~=0.4.25
 requests~=2.29.0
 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,111 @@
+import { useSelector } from '@xstate/react';
+import React, { useCallback, useState } from 'react';
+import { useEditSegment, useSelect, useRaw } from '../../../ProjectContext';
+import ActionButton from './ActionButton';
+import { MenuItem, TextField, Box } from '@mui/material';
+import Grid from '@mui/material/Grid';
+import AddIcon from '@mui/icons-material/Add';
+import Button from '@mui/material/Button';
+
+function LayerSelector({ layer, channelType }) {
+  const segment = useEditSegment();
+  const primaryChannel = useSelector(segment, (state) => state.context.primaryChannel);
+  const secondaryChannel = useSelector(segment, (state) => state.context.secondaryChannel);
+
+  const raw = useRaw();
+  const names = useSelector(raw, (state) => state.context.channelNames);
+
+  const onChangePrimary = (e) => {
+    segment.send({ type: 'SET_PRIMARY_CHANNEL', primaryChannel: Number(e.target.value) });
+  };
+
+  const onChangeSecondary = (e) => {
+    segment.send({ type: 'SET_SECONDARY_CHANNEL', secondaryChannel: Number(e.target.value) });
+  };
+
+  return channelType == 'primary' ? (
+    <TextField
+      select
+      size='small'
+      value={primaryChannel}
+      onChange={onChangePrimary}
+      sx={{ width: 130 }}
+    >
+      {names.map((opt, index) => (
+        <MenuItem key={index} value={index}>
+          {opt}
+        </MenuItem>
+      ))}
+    </TextField>
+  ) : (
+    <TextField
+      select
+      size='small'
+      value={secondaryChannel}
+      onChange={onChangeSecondary}
+      sx={{ width: 130 }}
+    >
+      {names.map((opt, index) => (
+        <MenuItem key={index} value={index}>
+          {opt}
+        </MenuItem>
+      ))}
+    </TextField>
+  );
+}
+
+function SegmentAllButton({ props, layer }) {
+  const segment = useEditSegment();
+
+  const onClick = useCallback(() => segment.send('SEGMENTALL'), [segment]);
+  const [addButtonClicked, setAddButtonClicked] = useState(false);
+  const [boxes, setBoxes] = useState([]);
+
+  const addChannel = () => {
+    setBoxes([...boxes, { layer: layer, channelType: 'secondary' }]);
+    setAddButtonClicked(true);
+  };
+
+  return (
+    <Grid>
+      Select Primary Channel
+      <Grid item xs={10.5}>
+        <LayerSelector layer={layer} channelType={'primary'} />
+      </Grid>
+      Select Secondary Channel
+      <Grid item xs={10.5}>
+        <Box sx={{ minWidth: 140, marginTop: 0.4 }}>
+          {!addButtonClicked && (
+            <Button
+              onClick={addChannel}
+              fullWidth
+              variant='outlined'
+              sx={{ borderStyle: 'dashed', p: 0.5 }}
+              startIcon={<AddIcon />}
+              size='small'
+            >
+              Secondary Channel
+            </Button>
+          )}
+        </Box>
+        {boxes.map((box, index) => (
+          <Box key={index} sx={{ minWidth: 140, marginTop: 0.4 }}>
+            <LayerSelector layer={box.layer} channelType={box.channelType} />
+          </Box>
+        ))}
+      </Grid>
+      <Grid>
+        <ActionButton
+          {...props}
+          tooltipText={'Run cell sam on selected channel to get segmentation masks'}
+          onClick={onClick}
+          hotkey='m'
+        >
+          Segment All
+        </ActionButton>
+      </Grid>
+    </Grid>
+  );
+}
+
+export default SegmentAllButton;

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

@@ -0,0 +1,12 @@
+import Box from '@mui/material/Box';
+import ActionButtons from './ActionButtons';
+
+function SegmentSamControls() {
+  return (
+    <Box display='flex' flexDirection='column'>
+      <ActionButtons />
+    </Box>
+  );
+}
+
+export default SegmentSamControls;