feat(app): support preview grounding annotations

IDEA-Research · Apr 9, 2024 · 6ee8da9 · 6ee8da9
1 parent 4285e5f
commit 6ee8da9
Show file tree

Hide file tree

Showing 8 changed files with 178 additions and 10 deletions.
diff --git a/packages/app/src/locales/en-US.ts b/packages/app/src/locales/en-US.ts
@@ -102,6 +102,8 @@ const localeTexts = {
   'dataset.detail.pagination': 'Pagination',
   'dataset.detail.random': 'Random',
   'dataset.detail.randomQuery': 'Random',
+  'dataset.detail.showGrounding': 'Show Grounding',
+  'dataset.detail.hideGrounding': 'Hide Grounding',
 
   'dataset.toAnalysis.unSupportWarn':
     'You should have a prediction label set with detection annotaion first',

diff --git a/packages/app/src/locales/zh-CN.ts b/packages/app/src/locales/zh-CN.ts
@@ -100,6 +100,8 @@ const localeTexts = {
   'dataset.toAnalysis.unSelectWarn': '请选择一个预标注集',
   'dataset.onClickCopyLink.success': '复制链接成功!',
   'dataset.detail.overlay': '覆盖',
+  'dataset.detail.showGrounding': '展示 Grounding',
+  'dataset.detail.hideGrounding': '隐藏 Grounding',
 
   'dataset.filter.newDataset': '新建数据集',
   'dataset.filter.public': '公共',

diff --git a/packages/components/src/Annotator/components/HighlightText/index.less b/packages/components/src/Annotator/components/HighlightText/index.less
@@ -0,0 +1,5 @@
+.dds-annotator-highlight-text .ant-tag {
+	margin-inline-end: 0px !important;
+  font-size: 14px;
+  cursor: pointer;
+}
diff --git a/packages/components/src/Annotator/components/HighlightText/index.tsx b/packages/components/src/Annotator/components/HighlightText/index.tsx
@@ -0,0 +1,73 @@
+import React, { useMemo } from 'react';
+import { escapeRegExp } from 'lodash';
+import { Tag } from 'antd';
+
+import './index.less';
+
+interface IHighlight {
+  text: string;
+  color: string;
+}
+
+interface IProps {
+  text: string;
+  highlights: IHighlight[];
+  onHoverHighlightWord: (text: string) => void;
+  onLeaveHighlightWord: () => void;
+}
+
+const HighlightText: React.FC<IProps> = ({
+  text,
+  highlights,
+  onHoverHighlightWord,
+  onLeaveHighlightWord,
+}) => {
+
+  const segments = useMemo(() => {
+    const computedSegments: React.ReactNode[] = [];
+    const pattern = new RegExp(
+      highlights.map(h => `\\b(${escapeRegExp(h.text)})\\b`).join('|'),
+      'g'
+    );
+
+    const matches = Array.from(text.matchAll(pattern));
+    let lastIndex = 0;
+
+    matches.forEach(match => {
+      const matchText = match[0];
+      const index = match.index ?? 0;
+
+      if (index > lastIndex) {
+        computedSegments.push(text.substring(lastIndex, index));
+      }
+
+      const highlightConfig = highlights.find(h => h.text === matchText);
+
+      if (highlightConfig) {
+        computedSegments.push(
+          <Tag
+            key={`${index}-${matchText}`}
+            color={highlightConfig.color}
+            bordered={false}
+            onMouseEnter={() => onHoverHighlightWord(matchText)}
+            onMouseLeave={onLeaveHighlightWord}
+          >
+            {matchText}
+          </Tag>
+        );
+      }
+
+      lastIndex = index + matchText.length;
+    });
+
+    if (lastIndex < text.length) {
+      computedSegments.push(text.substring(lastIndex));
+    }
+
+    return computedSegments;
+  }, [text, highlights, onHoverHighlightWord, onLeaveHighlightWord]);
+
+  return <div className={'dds-annotator-highlight-text'}>{segments}</div>;
+};
+
+export default HighlightText;
diff --git a/packages/components/src/Annotator/hooks/useCanvasRender.tsx b/packages/components/src/Annotator/hooks/useCanvasRender.tsx
@@ -167,7 +167,7 @@ const useCanvasRender = ({
       } else if (
         theDrawData.selectedTool === EBasicToolItem.Rectangle &&
         theDrawData.selectedModel[theDrawData.selectedTool] ===
-          EnumModelType.IVP
+        EnumModelType.IVP
       ) {
         objectHooksMap[EObjectType.Rectangle].renderPrompt({
           prompt,
@@ -213,8 +213,8 @@ const useCanvasRender = ({
     const status = isFocus
       ? 'focus'
       : isJustCreated
-      ? 'justCreated'
-      : undefined;
+        ? 'justCreated'
+        : undefined;
     const styles = getObjectStyles(object, object.color, status);
 
     // Change globalAlpha when creating / editing object
@@ -315,6 +315,20 @@ const useCanvasRender = ({
         false,
       );
     }
+
+    // render highlight object when hover caption
+    if (!!drawData.highlightCategory) {
+      const highlights = theDrawData.objectList
+          .filter(obj => obj.labelId === drawData.highlightCategory!.id);
+
+      highlights.forEach((obj) => {
+        renderObject(
+          obj,
+          true,
+          false
+        );
+      });
+    }
   };
 
   const renderPopoverMenu = () => {
@@ -327,7 +341,7 @@ const useCanvasRender = ({
     ) {
       const target =
         drawData.objectList[editState.focusObjectIndex].keypoints?.points?.[
-          editState.focusEleIndex
+        editState.focusEleIndex
         ];
       if (target) {
         return (

diff --git a/packages/components/src/Annotator/index.less b/packages/components/src/Annotator/index.less
@@ -204,7 +204,7 @@
 
   .switch {
     position: absolute;
-    bottom: 40px;
+    top: 25px;
     display: flex;
     align-items: center;
     justify-content: center;
@@ -331,6 +331,18 @@
       transform: rotate(180deg);
     }
   }
+
+  .dds-annotator-grounding-preview {
+    position: absolute;
+    bottom: 20px;
+    left: 50%;
+    transform: translate(-50%, 0);
+    background: rgba(0, 0, 0, 0.45);
+    color: #fff;
+    width: 70%;
+    padding: 10px;
+    border-radius: 10px;
+  }
 }
 
 .dds-annotator-view {

diff --git a/packages/components/src/Annotator/preview.tsx b/packages/components/src/Annotator/preview.tsx
@@ -1,5 +1,7 @@
 import {
   CloseOutlined,
+  EyeInvisibleOutlined,
+  EyeOutlined,
   LeftOutlined,
   RightOutlined,
   ZoomInOutlined,
@@ -42,6 +44,9 @@ import {
 } from './type';
 
 import './index.less';
+import HighlightText from './components/HighlightText';
+import { useLocale } from 'dds-utils/locale';
+
 
 export interface PreviewProps {
   isOldMode?: boolean; // is old dataset design mode
@@ -70,6 +75,8 @@ const Preview: React.FC<PreviewProps> = (props) => {
     displayOptionsResult,
   } = props;
 
+  const { localeText } = useLocale();
+
   const [annotations, setAnnotations] = useImmer<DrawObject[]>([]);
 
   const [editState, setEditState] = useImmer<EditState>(
@@ -101,7 +108,7 @@ const Preview: React.FC<PreviewProps> = (props) => {
     allowMove: editState.allowMove,
     isRequiring: editState.isRequiring,
     minPadding: {
-      top: 120,
+      top: 150,
       left: 300,
     },
     cursorSize: drawData.brushSize,
@@ -240,12 +247,25 @@ const Preview: React.FC<PreviewProps> = (props) => {
   // =================================================================================================================
 
   const [showInfo, setShowInfo] = useState(true);
+  const [showGrounding, setShowGrounding] = useState(false);
+
+  const metadata = !isEmpty(list[current]?.metadata)
+    ? list[current].metadata
+    : undefined;
+
   const changeShowInfo = useCallback(() => {
     setShowInfo((s) => {
       return !s;
     });
   }, []);
 
+  const changeShowGrounding = useCallback(() => {
+    if (!metadata || !metadata.caption) return;
+    setShowGrounding((s) => {
+      return !s;
+    });
+  }, [metadata]);
+
   /** Snapshot image */
   const onDownload: React.MouseEventHandler<HTMLDivElement> = async (event) => {
     event.preventDefault();
@@ -318,7 +338,7 @@ const Preview: React.FC<PreviewProps> = (props) => {
     ) {
       const target =
         drawData.objectList[editState.focusObjectIndex]?.keypoints?.points?.[
-          editState.focusEleIndex
+        editState.focusEleIndex
         ];
       if (target) {
         return (
@@ -337,9 +357,30 @@ const Preview: React.FC<PreviewProps> = (props) => {
     return <></>;
   }
 
-  const metadata = !isEmpty(list[current]?.metadata)
-    ? list[current].metadata
-    : undefined;
+  const getHighlightWords = () => {
+    const objects = list[current].objects;
+    return categories
+      .filter((category) => objects.find((obj: any) => obj.categoryId === category.id))
+      .map((category) => {
+        return {
+          text: category.name,
+          color: getAnnotColor(category.id),
+        }
+      })
+  };
+
+  const highlightCategory = (name: string) => {
+    setDrawData((s) => {
+      const category = categories.find(item => item.name === name);
+      s.highlightCategory = category;
+    });
+  };
+
+  const clearHighlightCategory = () => {
+    setDrawData((s) => {
+      s.highlightCategory = undefined;
+    });
+  };
 
   return (
     <div className="dds-annotator dds-annotator-preview">
@@ -360,6 +401,12 @@ const Preview: React.FC<PreviewProps> = (props) => {
             icon: <DownloadIcon />,
             onClick: onDownload,
           },
+          {
+            icon: showGrounding ? <EyeOutlined /> : <EyeInvisibleOutlined />,
+            onClick: changeShowGrounding,
+            disabled: !metadata || !metadata.caption,
+            title: showGrounding ? localeText('dataset.detail.hideGrounding') : localeText('dataset.detail.showGrounding'),
+          },
         ]}
         rightTools={[
           {
@@ -431,6 +478,17 @@ const Preview: React.FC<PreviewProps> = (props) => {
           <DoubleRightIcon />
         </div>
       )}
+      {
+        showGrounding && !!metadata?.caption &&
+        <div className="dds-annotator-grounding-preview">
+          <HighlightText
+            text={metadata.caption}
+            highlights={getHighlightWords()}
+            onHoverHighlightWord={(text) => highlightCategory(text)}
+            onLeaveHighlightWord={clearHighlightCategory}
+          />
+        </div>
+      }
     </div>
   );
 };

diff --git a/packages/components/src/Annotator/type.ts b/packages/components/src/Annotator/type.ts
@@ -225,6 +225,7 @@ export interface DrawData {
   isBatchEditing: boolean; // active while handle batch predictions by model
   editingAttribute?: IEditingAttribute;
   limitConf: number;
+  highlightCategory?: Category;
 
   /** prompt actions */
   prompt: IPrompt;
@@ -296,6 +297,7 @@ export const DEFAULT_DRAW_DATA: DrawData = {
   isJustCreated: false,
   creatingObject: undefined,
   editingAttribute: undefined,
+  highlightCategory: undefined,
   brushSize: 20,
   pointResolution: 0.5,
   prompt: {},