feat: allow string as label in bbox lens

Renumics · Feb 19, 2024 · ae3b4c2 · ae3b4c2
1 parent 6438101
commit ae3b4c2
Showing 1 changed file with 48 additions and 40 deletions.
diff --git a/src/lenses/BoundingBoxLens/index.tsx b/src/lenses/BoundingBoxLens/index.tsx
@@ -1,6 +1,6 @@
 import { useColorTransferFunction } from '../../hooks';
 import { useRef, useCallback, useState, useMemo } from 'react';
-import { Lens } from '../../types';
+import { DataColumn, Lens } from '../../types';
 import 'twin.macro';
 import { CategoricalDataType } from '../../datatypes';
 import BBox from './BBox';
@@ -130,59 +130,67 @@ const BoundingBoxLens: Lens = ({ urls, values, columns }) => {
 };
 
 BoundingBoxLens.key = 'BoundingBoxView';
-BoundingBoxLens.dataTypes = ['Image', 'BoundingBox', 'Sequence', 'Category'];
+BoundingBoxLens.dataTypes = ['Image', 'BoundingBox', 'Sequence', 'Category', 'str'];
 BoundingBoxLens.multi = true;
 BoundingBoxLens.displayName = 'BoundingBox';
 BoundingBoxLens.defaultHeight = 256;
 BoundingBoxLens.filterAllowedColumns = (allColumns, selectedColumns) => {
-    if (selectedColumns.length === 3) return [];
-
     const selectedTypes = selectedColumns.map((selectedCol) => selectedCol.type);
-    if (selectedColumns.length === 0)
-        return allColumns.filter(({ type }) =>
-            BoundingBoxLens.dataTypes.includes(type.kind)
-        );
+    const allowedColumns: DataColumn[] = [];
 
-    if (selectedTypes.find((type) => type.kind === 'Sequence')) {
-        return allColumns.filter(({ type }) => {
-            return type.kind === 'Sequence' || type.kind === 'Image';
-        });
-    } else if (
-        selectedTypes.find((type) => ['BoundingBox', 'Category'].includes(type.kind))
-    ) {
-        return allColumns.filter(({ type }) => {
-            return (
+    // allow exactly one image column
+    if (!selectedTypes.find((type) => type.kind === 'Image')) {
+        allowedColumns.push(...allColumns.filter(({ type }) => type.kind === 'Image'));
+    }
+
+    // allow exactly one bbox or bbox[] column
+    if (
+        !selectedTypes.find(
+            (type) =>
                 type.kind === 'BoundingBox' ||
-                type.kind === 'Category' ||
-                type.kind === 'Image'
-            );
-        });
-    } else {
-        return allColumns.filter(({ type }) =>
-            BoundingBoxLens.dataTypes.includes(type.kind)
+                (type.kind === 'Sequence' && type.dtype.kind === 'BoundingBox')
+        )
+    ) {
+        allowedColumns.push(
+            ...allColumns.filter(
+                ({ type }) =>
+                    type.kind === 'BoundingBox' ||
+                    (type.kind === 'Sequence' && type.dtype.kind === 'BoundingBox')
+            )
         );
     }
-};
-BoundingBoxLens.isSatisfied = (columns) => {
-    if (columns.length < 2) return false;
-
-    const types = columns.map((col) => col.type);
-
-    if (!types.find((type) => type.kind === 'Image')) return false;
-
-    if (!types.find((type) => ['BoundingBox', 'Sequence'].includes(type.kind)))
-        return false;
 
+    // allow exactly one label or label[] column
     if (
-        types.find(
-            (type) => type.kind === 'Sequence' && type.dtype.kind === 'BoundingBox'
+        !selectedTypes.find(
+            (type) =>
+                type.kind === 'Category' ||
+                type.kind === 'str' ||
+                (type.kind === 'Sequence' &&
+                    (type.dtype.kind === 'str' || type.dtype.kind === 'Category'))
         )
     ) {
-        return true;
-    } else {
-        if (types.find((type) => type.kind === 'BoundingBox')) return true;
-        else return false;
+        allowedColumns.push(
+            ...allColumns.filter(
+                ({ type }) =>
+                    type.kind === 'Category' ||
+                    type.kind === 'str' ||
+                    (type.kind === 'Sequence' &&
+                        (type.dtype.kind === 'str' || type.dtype.kind === 'Category'))
+            )
+        );
     }
+    return allowedColumns;
+};
+BoundingBoxLens.isSatisfied = (columns) => {
+    const types = columns.map((col) => col.type);
+    const hasImage = !!types.find((type) => type.kind === 'Image');
+    const hasBbox = !!types.find(
+        (type) =>
+            type.kind === 'BoundingBox' ||
+            (type.kind === 'Sequence' && type.dtype.kind === 'BoundingBox')
+    );
+    return hasImage && hasBbox;
 };
 
 export default BoundingBoxLens;