[web] 250 mask annotation optimize

ymir/web/src/components/dataset/ListAnnotation.tsx

@@ -9,7 +9,7 @@ import BoundingBox from './asset/BoundingBox'
 
 import styles from './common.less'
 import { useDebounceEffect } from 'ahooks'
-import { Annotation, Asset } from '@/constants'
+import { Annotation, Asset, BoundingBox as BoundingBoxType, Polygon as PolygonType, Mask as MaskType } from '@/constants'
 
 type Props = {
   asset: Asset
@@ -26,6 +26,19 @@ const ListAnnotation: FC<Props> = ({ asset, filter, hideAsset, isFull }) => {
   const [imgWidth, setImgWidth] = useState(100)
   const [ratio, setRatio] = useState(1)
 
+  const [bbA, setBbA] = useState<BoundingBoxType[]>([])
+  const [pgA, setPgA] = useState<PolygonType[]>([])
+  const [maA, setMaA] = useState<MaskType[]>([])
+
+  useEffect(() => {
+    if (!annotations?.length) {
+      return
+    }
+    setBbA(filterAts(AnnotationType.BoundingBox) as BoundingBoxType[])
+    setPgA(filterAts(AnnotationType.Polygon) as PolygonType[])
+    setMaA(filterAts(AnnotationType.Mask) as MaskType[])
+  }, [annotations])
+
   useDebounceEffect(
     () => {
       if (!asset) {
@@ -50,22 +63,15 @@ const ListAnnotation: FC<Props> = ({ asset, filter, hideAsset, isFull }) => {
 
   window.addEventListener('resize', () => imgContainer.current && calClientWidth())
 
-  function renderAnnotation(annotation: Annotation) {
-    switch (annotation.type) {
-      case AnnotationType.BoundingBox:
-        return <BoundingBox key={annotation.id} annotation={annotation} ratio={ratio} simple={true} />
-      case AnnotationType.Polygon:
-        return <Polygon key={annotation.id} annotation={annotation} ratio={ratio} simple={true} />
-      case AnnotationType.Mask:
-        return <Mask key={annotation.id} annotation={annotation} ratio={ratio} simple={true} />
-    }
-  }
+  const filterAts = (fType: AnnotationType) => annotations?.filter(({ type }) => fType === type)
 
   return (
-    <div className={styles.ic_container} ref={imgContainer} key={asset.hash}>
+    <div className={styles.ic_container} ref={imgContainer}>
       <img ref={img} style={{ visibility: hideAsset ? 'hidden' : 'visible' }} src={asset?.url} className={styles.assetImg} onLoad={calClientWidth} />
       <div className={styles.annotations} style={{ width: imgWidth, left: -imgWidth / 2 }}>
-        {annotations.map(renderAnnotation)}
+        {bbA.length ? bbA.map((anno) => <BoundingBox key={anno.id} annotation={anno} ratio={ratio} simple={true} />) : null}
+        {pgA.length ? <Polygon annotations={pgA} width={asset.width} height={asset.height} ratio={ratio} simple={true} /> : null}
+        {maA.length ? maA.map((anno) => <Mask key={anno.id} annotation={anno} width={asset.width} height={asset.height} ratio={ratio} simple={true} />) : null}
       </div>
     </div>
   )

ymir/web/src/components/dataset/add/Local.tsx

@@ -3,12 +3,12 @@ import { useSelector } from 'umi'
 import { Types } from './AddTypes'
 import { Button, Form } from 'antd'
 import useRequest from '@/hooks/useRequest'
+import { formLayout } from '@/config/antd'
+import SubmitBtn from './SubmitBtn'
 import { ImportingItem } from '@/constants'
 import t from '@/utils/t'
 import Uploader from '@/components/form/uploader'
 import Tip from './Tip'
-import { formLayout } from '@/config/antd'
-import SubmitBtn from './SubmitBtn'
 
 const Local: FC = () => {
   const [form] = Form.useForm()

ymir/web/src/components/dataset/asset/AssetAnnotations.tsx

@@ -8,7 +8,7 @@ import Polygon from './Polygon'
 import BoundingBox from './BoundingBox'
 
 import styles from '../common.less'
-import { Annotation } from '@/constants'
+import { Annotation, BoundingBox as BoundingBoxType, Mask as MaskType, Polygon as PolygonType } from '@/constants'
 
 export type Asset = {
   annotations: Annotation[]
@@ -29,6 +29,19 @@ const AssetAnnotation: FC<Props> = ({ asset }) => {
   const [imgWidth, setImgWidth] = useState(0)
   const [ratio, setRatio] = useState(1)
 
+  const [bbA, setBbA] = useState<BoundingBoxType[]>([])
+  const [pgA, setPgA] = useState<PolygonType[]>([])
+  const [maA, setMaA] = useState<MaskType[]>([])
+
+  useEffect(() => {
+    if (!annotations?.length) {
+      return
+    }
+    setBbA(filterAts(AnnotationType.BoundingBox) as BoundingBoxType[])
+    setPgA(filterAts(AnnotationType.Polygon) as PolygonType[])
+    setMaA(filterAts(AnnotationType.Mask) as MaskType[])
+  }, [annotations])
+
   useEffect(() => {
     if (!asset) {
       return
@@ -55,17 +68,6 @@ const AssetAnnotation: FC<Props> = ({ asset }) => {
     setRatio(clientWidth / iw)
   }
 
-  function renderAnnotation(annotation: Annotation) {
-    switch (annotation.type) {
-      case AnnotationType.BoundingBox:
-        return <BoundingBox key={annotation.id} annotation={annotation} ratio={ratio} />
-      case AnnotationType.Polygon:
-        return <Polygon key={annotation.id} annotation={annotation} ratio={ratio} />
-      case AnnotationType.Mask:
-        return <Mask key={annotation.id} annotation={annotation} ratio={ratio} />
-    }
-  }
-
   const imgLoad = (e: SyntheticEvent) => {
     if (img.current && img.current.naturalWidth) {
       const { naturalWidth } = img.current
@@ -79,13 +81,17 @@ const AssetAnnotation: FC<Props> = ({ asset }) => {
     }
   })
 
+  const filterAts = (fType: AnnotationType) => annotations?.filter(({ type }) => fType === type)
+
   return (
     <div className={styles.anno_panel} ref={imgContainer} key={asset?.hash}>
       <div className={styles.img_container}>
         <img ref={img} src={asset?.url} style={imgWidth ? { width: imgWidth } : undefined} className={styles.assetImg} onLoad={imgLoad} />
       </div>
       <div className={styles.annotations} style={{ width: imgWidth, left: -imgWidth / 2 }}>
-        {annotations.map(renderAnnotation)}
+        {bbA.length ? bbA.map((anno) => <BoundingBox key={anno.id} annotation={anno} ratio={ratio} />) : null}
+        {pgA.length ? <Polygon annotations={pgA} width={asset.width} height={asset.height} ratio={ratio} /> : null}
+        {maA.length ? maA.map((anno) => <Mask key={anno.id} annotation={anno} width={asset.width} height={asset.height} ratio={ratio} />) : null}
       </div>
     </div>
   )

ymir/web/src/components/dataset/asset/Mask.tsx

@@ -6,37 +6,25 @@ type Props = {
   annotation: MaskType
   ratio?: number
   simple?: boolean
+  width?: number
+  height?: number
 }
 
-const Mask: FC<Props> = ({ annotation, ratio = 1 }) => {
+const Mask: FC<Props> = ({ annotation, ratio = 1, width, height, simple }) => {
   const canvasRef = useRef<HTMLCanvasElement>(null)
   const [canvas, setCanvas] = useState<HTMLCanvasElement>()
-  const [{ width, height }, setRect] = useState({
-    width: 0,
-    height: 0,
-  })
+
   useEffect(() => {
     if (canvasRef.current) {
       setCanvas(canvasRef.current)
     }
   }, [canvasRef.current])
 
   useEffect(() => {
-    if (annotation.decodeMask && canvas) {
-      const { decodeMask: mask, color } = annotation
-      renderMask(canvas, mask, width, height, color)
-    }
-  }, [annotation.decodeMask, canvas, width, height])
-
-  useEffect(() => {
-    if (!annotation) {
-      return
+    if (annotation && canvas && width && height) {
+      renderMask(canvas, annotation, !simple, ratio)
     }
-    setRect({
-      width: annotation.width,
-      height: annotation.height,
-    })
-  }, [annotation])
+  }, [annotation, canvas, width, height])
 
   return (
     <canvas

ymir/web/src/components/dataset/asset/Polygon.tsx

@@ -1,40 +1,29 @@
 import { FC, useEffect, useRef, useState } from 'react'
-import { renderPolygon } from './_helper'
+import { renderPolygons } from './_helper'
 import { Polygon as PolygonType } from '@/constants'
+
 type Props = {
-  annotation: PolygonType
+  annotations: PolygonType[]
   ratio?: number
   simple?: boolean
+  width?: number
+  height?: number
 }
 
-const Polygon: FC<Props> = ({ annotation, ratio = 1 }) => {
+const Polygon: FC<Props> = ({ annotations, ratio = 1, width, height, simple }) => {
   const canvasRef = useRef<HTMLCanvasElement>(null)
   const [canvas, setCanvas] = useState<HTMLCanvasElement>()
-  const [{ width, height }, setRect] = useState({
-    width: 0,
-    height: 0,
-  })
   useEffect(() => {
     if (canvasRef.current) {
       setCanvas(canvasRef.current)
     }
   }, [canvasRef.current])
 
   useEffect(() => {
-    if (annotation.polygon?.length && canvas) {
-      renderPolygon(canvas, annotation.polygon, annotation.color)
-    }
-  }, [annotation.polygon, canvas, width, height])
-
-  useEffect(() => {
-    if (!annotation) {
-      return
+    if (annotations.length && canvas && width && height) {
+      renderPolygons(canvas, annotations, !simple, ratio)
     }
-    setRect({
-      width: annotation.width,
-      height: annotation.height,
-    })
-  }, [annotation])
+  }, [annotations, canvas, width, height])
 
   return (
     <canvas

ymir/web/src/components/dataset/asset/_helper.ts

@@ -29,28 +29,59 @@ function mask2Uint8Array(mask: number[][], len: number, color?: string) {
   return dataWithColor
 }
 
-export function renderPolygon(canvas: HTMLCanvasElement, points: Polygon['polygon'], color?: string) {
+export function renderPolygons(canvas: HTMLCanvasElement, annotations: Polygon[], showMore?: boolean, ratio?: number) {
   const ctx = canvas.getContext('2d')
-  if (!ctx) {
+  if (!ctx || !annotations.length) {
     return
   }
-  ctx.beginPath()
+  const { color } = annotations[0]
   ctx.fillStyle = getColor(color).hexa()
-  ctx.strokeStyle = getColor('black').hexa()
-  ctx.moveTo(points[0].x, points[0].y)
+  ctx.strokeStyle = getColor('gray').hexa()
   ctx.lineWidth = 1
-  points.forEach((point, index) => index > 0 && ctx.lineTo(point.x, point.y))
+  ctx.beginPath()
+  annotations.forEach((annotation) => {
+    const { polygon: points } = annotation
+    ctx.moveTo(points[0].x, points[0].y)
+    points.forEach((point, index) => index > 0 && ctx.lineTo(point.x, point.y))
+  })
   ctx.fill()
+  showMore && drawBoxs(ctx, annotations, ratio)
 }
 
-export function renderMask(canvas: HTMLCanvasElement, mask: number[][], width: number, height: number, color?: string) {
+export function renderMask(canvas: HTMLCanvasElement, annotation: Mask, showMore?: boolean, ratio?: number) {
   const ctx = canvas.getContext('2d')
   if (!ctx) {
     return
   }
-  const image = mask2Image(mask, width, height, color)
+  const { width, height, color, decodeMask } = annotation
+  if (!decodeMask) {
+    return
+  }
+  const image = mask2Image(decodeMask, width, height, color)
 
   image && ctx.putImageData(image, 0, 0)
+
+  showMore && drawBoxs(ctx, [annotation], ratio)
+}
+
+function drawBoxs(ctx: CanvasRenderingContext2D, annotations: Annotation[], ratio: number = 1) {
+  const { color } = annotations[0]
+  const lw = 1 / ratio
+  const th = 16 * lw
+  const padding = lw * 4
+  const mainColor = getColor(color).hexa()
+  ctx.strokeStyle = mainColor
+  ctx.lineWidth = 1 / ratio
+  ctx.font = `${th}px Microsoft Yahei`
+  annotations.forEach(({ box, keyword, score }) => {
+    ctx.strokeRect(box.x, box.y, box.w, box.h)
+    const text = `${keyword} ${score || ''}`
+    const tw = ctx.measureText(text).width
+    ctx.fillStyle = getColor(color, 0.6).hexa()
+    ctx.fillRect(box.x, box.y - th - padding * 2, tw + 8, th + padding * 2)
+    ctx.fillStyle = 'white'
+    ctx.fillText(text, box.x + padding, box.y - padding)
+  })
 }
 
 export function transferAnnotations(annotations: Annotation[] = []) {

ymir/web/src/constants/typings/asset.d.ts

@@ -35,6 +35,7 @@ interface AnnotationBase {
   score?: number
   gt?: boolean
   cm: number
+  box: Bbox
   tags?: CustomClass
 }
 
@@ -54,15 +55,15 @@ type Annotation = Matable<AnnotationMaps>
 type SegAnnotation = Matable<Omit<AnnotationMaps, AnnotationType.BoundingBox>>
 type DetAnnotation = Matable<Pick<AnnotationMaps, AnnotationType.BoundingBox>>
 
+type Bbox = {
+  x: number
+  y: number
+  w: number
+  h: number
+  rotate_angle?: number
+}
 interface BoundingBox extends AnnotationBase {
   type: AnnotationType.BoundingBox
-  box: {
-    x: number
-    y: number
-    w: number
-    h: number
-    rotate_angle?: number
-  }
 }
 
 interface Polygon extends AnnotationBase {
@@ -74,7 +75,6 @@ interface Mask extends AnnotationBase {
   type: AnnotationType.Mask
   mask: string
   decodeMask?: number[][]
-  rect?: [x: number, y: number, width: number, height: number]
 }
 
 export { Asset, AnnotationBase, Annotation, SegAnnotation, DetAnnotation, BoundingBox, Polygon, Mask }