Change raw rescaling to happen on frontend + export original raw unsc…

…aled (#383)

* Changed export to give original raw input

* Change REPLACE tool name with COMBINE:

* Apply automatic changes

* Correctly get raw dtype on export

* Change matching to mapping for python <3.10

* Apply automatic changes

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

.gitignore

@@ -150,6 +150,10 @@ fmd_config.cfg
 # production
 /build
 
+# models and other pickled files
+*.h5
+*.npy
+
 # misc
 .DS_Store
 .env.local

backend/deepcell_label/export.py

@@ -36,6 +36,22 @@ def load_dimensions(self):
                 self.duration = dimensions['duration']
                 self.num_channels = dimensions['numChannels']
                 self.num_features = dimensions['numFeatures']
+                self.dtype = self.get_dtype(dimensions['dtype'])
+
+    def get_dtype(self, arr_type):
+        """Matches raw array dtype with a numpy dtype"""
+        mapping = {
+            'Uint8Array': np.uint8,
+            'Uint16Array': np.uint16,
+            'Uint32Array': np.uint32,
+            'Int32Array': np.int32,
+            'Float32Array': np.float32,
+            'Float64Array': np.float64,
+        }
+        try:
+            return mapping[arr_type]
+        except KeyError:
+            raise ValueError('Could not match dtype of raw array.')
 
     def load_labeled(self):
         """Loads the labeled array from the labeled.dat file."""
@@ -51,7 +67,8 @@ def load_raw(self):
         """Loads the raw array from the raw.dat file."""
         with zipfile.ZipFile(self.labels_zip) as zf:
             with zf.open('raw.dat') as f:
-                raw = np.frombuffer(f.read(), np.uint8)
+                raw = np.frombuffer(f.read(), self.dtype)
+                print(raw.dtype)
                 self.raw = np.reshape(
                     raw, (self.num_channels, self.duration, self.height, self.width)
                 )

backend/deepcell_label/loaders.py

@@ -78,10 +78,6 @@ def write_images(self):
         """
         X = self.X
         if X is not None:
-            # Rescale data
-            max, min = np.max(X), np.min(X)
-            X = (X - min) / (max - min if max - min > 0 else 1) * 255
-            X = X.astype(np.uint8)
             # Move channel axis
             X = np.moveaxis(X, -1, 1)
             images = io.BytesIO()

frontend/src/Project/EditControls/CellControls/ToolButtons.js

@@ -47,7 +47,7 @@ function ToolButtons() {
         </ToggleButton>
         <ToggleButton value={'replace'} sx={{ px: 0.5, py: 0 }}>
           <Tooltip title={<kbd>R</kbd>} placement='right'>
-            <div>Replace</div>
+            <div>Combine</div>
           </Tooltip>
         </ToggleButton>
         <ToggleButton value={'swap'} sx={{ px: 0.5, py: 0 }}>

frontend/src/Project/Instructions/CellsInstructions.js

@@ -10,7 +10,7 @@ function ToolShortcuts() {
     <Shortcuts>
       <Shortcut text='Select' shortcut='V' />
       <Shortcut text='Delete' shortcut='Backspace' />
-      <Shortcut text='Replace' shortcut='R' />
+      <Shortcut text='Combine' shortcut='R' />
       <Shortcut text='Swap' shortcut='S' />
       <Shortcut text='New' shortcut='N' />
     </Shortcuts>
@@ -65,7 +65,7 @@ function CellsInstructions() {
               <br />
               <strong>Delete</strong> removes the cell.
               <br />
-              <strong>Replace</strong> combines a second cell with the selected cell. First selects
+              <strong>Combine</strong> combines a second cell with the selected cell. First selects
               a cell, and then selects the cell to replace.
               <br />
               <strong>Swap</strong> switches a second cell with the selected cell. First picks

frontend/src/Project/Instructions/DivisionsInstructions.js

@@ -36,7 +36,7 @@ function DivisionsInstructions() {
             <strong>Dismiss:</strong> If the division is on the edge, the second daughter may be off
             screen.
             <br />
-            <strong>Fix:</strong> Combine the parent and daughter with Replace in the Cells tab.
+            <strong>Fix:</strong> Combine the parent and daughter with Combine in the Cells tab.
           </Typography>
         </Grid>
         <Grid item xs={3}>

frontend/src/Project/service/exportMachine.js

@@ -14,6 +14,7 @@ async function makeExportZip(context) {
     duration: raw[0].length,
     numChannels: raw.length,
     numFeatures: labeled.length,
+    dtype: raw[0][0][0].constructor.name,
   };
   const zipWriter = new zip.ZipWriter(new zip.BlobWriter('application/zip'));
   await zipWriter.add('dimensions.json', new zip.TextReader(JSON.stringify(dimensions)));

frontend/src/Project/service/labels/arraysMachine.js

@@ -24,6 +24,7 @@ const createArraysMachine = (context) =>
       context: {
         raw: null,
         labeled: null,
+        rawOriginal: null,
         t: 0,
         feature: 0,
         channel: 0,
@@ -48,7 +49,7 @@ const createArraysMachine = (context) =>
               states: {
                 loading: {
                   on: {
-                    LOADED: { target: 'done', actions: ['setRaw', 'setLabeled'] },
+                    LOADED: { target: 'done', actions: ['setRaw', 'setLabeled', 'setRawOriginal'] },
                   },
                 },
                 done: { type: 'final' },
@@ -121,6 +122,7 @@ const createArraysMachine = (context) =>
       actions: {
         setRaw: assign({ raw: (ctx, evt) => evt.raw }),
         setLabeled: assign({ labeled: (ctx, evt) => evt.labeled }),
+        setRawOriginal: assign({ rawOriginal: (ctx, evt) => evt.rawOriginal }),
         setT: assign({ t: (ctx, evt) => evt.t }),
         setFeature: assign({ feature: (ctx, evt) => evt.feature }),
         setChannel: assign({ channel: (ctx, evt) => evt.channel }),
@@ -148,7 +150,7 @@ const createArraysMachine = (context) =>
         ),
         sendArrays: respond((ctx) => ({
           type: 'ARRAYS',
-          raw: ctx.raw,
+          raw: ctx.rawOriginal, // Send the original raw for export
           labeled: ctx.labeled,
         })),
         save: send('SAVE', { to: (ctx) => ctx.undoRef }),

frontend/src/Project/service/loadMachine.ts

@@ -72,10 +72,11 @@ async function splitArrays(files: Files) {
   const labeledFile = files['y.ome.tiff'] as OmeTiff;
   const raw = await getRawRasters(rawFile.data[0]);
   const labeled = await getLabelRasters(labeledFile.data[0]);
-  return { raw, labeled };
+  const rawOriginal = await getRawOriginalRasters(rawFile.data[0]);
+  return { raw, labeled, rawOriginal };
 }
 
-async function getRawRasters(source: TiffPixelSource) {
+async function getRawOriginalRasters(source: TiffPixelSource) {
   const { labels, shape } = source;
   const c = shape[labels.indexOf('c')];
   const z = shape[labels.indexOf('z')];
@@ -85,14 +86,46 @@ async function getRawRasters(source: TiffPixelSource) {
     for (let j = 0; j < z; j++) {
       const selection = { t: 0, c: i, z: j };
       const raster = (await source.getRaster({ selection })) as Raster;
-      const frame = splitRows(raster) as Uint8Array[];
+      const frame = splitRows(raster);
       frames.push(frame);
     }
     channels.push(frames);
   }
   return channels;
 }
 
+async function getRawRasters(source: TiffPixelSource) {
+  const { labels, shape } = source;
+  const c = shape[labels.indexOf('c')];
+  const z = shape[labels.indexOf('z')];
+  const channels = [];
+  var max = 0;
+  var min = Infinity;
+  for (let i = 0; i < c; i++) {
+    const frames = [];
+    for (let j = 0; j < z; j++) {
+      const selection = { t: 0, c: i, z: j };
+      const raster = (await source.getRaster({ selection })) as Raster;
+      const frame = splitRows(raster);
+      // Record max and min across frames
+      for (let k = 0; k < frame.length; k++) {
+        for (let l = 0; l < frame[k].length; l++) {
+          if (frame[k][l] > max) {
+            max = frame[k][l];
+          }
+          if (frame[k][l] < min) {
+            min = frame[k][l];
+          }
+        }
+      }
+      frames.push(frame);
+    }
+    channels.push(frames);
+  }
+  const reshaped = reshapeRaw(channels, min, max) as Uint8Array[][][];
+  return reshaped;
+}
+
 async function getLabelRasters(source: TiffPixelSource) {
   const { labels, shape } = source;
   const c = shape[labels.indexOf('c')];
@@ -111,15 +144,47 @@ async function getLabelRasters(source: TiffPixelSource) {
   return channels;
 }
 
-type Raster = { data: Uint8Array | Int32Array; width: number; height: number };
+function reshapeRaw(channels: TypedArray[][][], min: number, max: number) {
+  // Normalize each pixel to 0-255 for rendering
+  for (let c = 0; c < channels.length; c++) {
+    for (let z = 0; z < channels[c].length; z++) {
+      for (let y = 0; y < channels[c][z].length; y++) {
+        for (let x = 0; x < channels[c][z][y].length; x++) {
+          channels[c][z][y][x] = Math.round(((channels[c][z][y][x] - min) / (max - min)) * 255);
+        }
+      }
+    }
+  }
+  return channels;
+}
+
+type TypedArray =
+  | Uint8Array
+  | Int8Array
+  | Uint16Array
+  | Int16Array
+  | Uint32Array
+  | Int32Array
+  | Float32Array
+  | Float64Array;
+
+type Raster = { data: TypedArray; width: number; height: number };
 
 function splitRows(raster: Raster) {
   const { data, width, height } = raster;
   const frame = [];
   for (let i = 0; i < height; i++) {
     const row =
-      data instanceof Uint8Array
+      data instanceof Uint8Array || data instanceof Int8Array
         ? new Uint8Array(data.buffer, width * i, width)
+        : data instanceof Uint16Array || data instanceof Int16Array
+        ? new Uint16Array(data.buffer, width * i * 2, width)
+        : data instanceof Uint32Array
+        ? new Uint32Array(data.buffer, width * i * 4, width)
+        : data instanceof Float32Array
+        ? new Float32Array(data.buffer, width * i * 4, width)
+        : data instanceof Float64Array
+        ? new Float64Array(data.buffer, width * i * 8, width)
         : new Int32Array(data.buffer, width * i * 4, width);
     frame.push(row);
   }
@@ -137,6 +202,7 @@ interface Context {
   numFeatures: number | null;
   raw: Uint8Array[][][] | null;
   labeled: Int32Array[][][] | null;
+  rawOriginal: TypedArray[][][] | null;
   labels: Cells | null;
   spots: Spots | null;
   divisions: Divisions | null;
@@ -157,6 +223,7 @@ const createLoadMachine = (projectId: string) =>
         numFeatures: null,
         raw: null,
         labeled: null,
+        rawOriginal: null,
         labels: null,
         spots: null,
         divisions: null,
@@ -175,6 +242,7 @@ const createLoadMachine = (projectId: string) =>
             data: {
               raw: Uint8Array[][][];
               labeled: Int32Array[][][];
+              rawOriginal: TypedArray[][][];
             };
           };
         },
@@ -244,11 +312,13 @@ const createLoadMachine = (projectId: string) =>
         'set arrays': assign({
           raw: (_, event) => event.data.raw,
           labeled: (_, event) => event.data.labeled,
+          rawOriginal: (_, event) => event.data.rawOriginal,
         }),
         'send loaded': sendParent((ctx) => ({
           type: 'LOADED',
           raw: ctx.raw,
           labeled: ctx.labeled,
+          rawOriginal: ctx.rawOriginal,
           spots: ctx.spots,
           divisions: ctx.divisions,
           cells: ctx.cells,

frontend/src/Project/service/projectMachine.js

@@ -103,7 +103,7 @@ const createProjectMachine = (projectId) =>
         }),
         sendDimensions: send(
           (c, e) => {
-            const { raw, labeled } = e;
+            const { raw, labeled, rawOriginal } = e;
             return {
               type: 'DIMENSIONS',
               numChannels: raw.length,