feat: Aux points support. #1597

h2oai · Nov 3, 2022 · 85c8436 · 85c8436
1 parent f1c03ce
commit 85c8436
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 24 deletions.
diff --git a/ui/src/image_annotator.test.tsx b/ui/src/image_annotator.test.tsx
@@ -343,6 +343,18 @@ describe('ImageAnnotator.tsx', () => {
       expect(wave.args[name]).toMatchObject([rect])
     })
 
+    it('Adds aux point to polygon when clicked', async () => {
+      const { container } = render(<XImageAnnotator model={model} />)
+      await waitForLoad()
+      const canvasEl = container.querySelectorAll('canvas')[1]
+      fireEvent.click(canvasEl, { clientX: 180, clientY: 120 })
+      fireEvent.click(canvasEl, { clientX: 240, clientY: 160 })
+      expect(wave.args[name]).toMatchObject([
+        rect,
+        { shape: { polygon: { items: [{ x: 105, y: 100 }, { x: 240, y: 100 }, { x: 240, y: 160 }, { x: 240, y: 220 },] } }, tag: 'person' }
+      ])
+    })
+
     it('Changes tag of existing polygon when clicked ', async () => {
       const { container, getByText } = render(<XImageAnnotator model={model} />)
       await waitForLoad()
@@ -371,6 +383,16 @@ describe('ImageAnnotator.tsx', () => {
       expect(canvasEl.style.cursor).toBe('auto')
     })
 
+    it('Displays the correct cursor when hovering over polygon aux point', async () => {
+      const { container } = render(<XImageAnnotator model={model} />)
+      await waitForLoad()
+      const canvasEl = container.querySelectorAll('canvas')[1]
+      fireEvent.click(canvasEl, { clientX: 180, clientY: 120 })
+      expect(canvasEl.style.cursor).toBe('move')
+      fireEvent.mouseMove(canvasEl, { clientX: 240, clientY: 160 })
+      expect(canvasEl.style.cursor).toBe('pointer')
+    })
+
     it('Moves polygon correctly', async () => {
       const { container } = render(<XImageAnnotator model={model} />)
       await waitForLoad()

diff --git a/ui/src/image_annotator.tsx b/ui/src/image_annotator.tsx
@@ -123,7 +123,7 @@ const
         : 'crosshair'
     if (intersected?.isFocused && intersected.shape.rect) cursor = getRectCornerCursor(intersected.shape.rect, cursor_x, cursor_y) || 'move'
     else if (focused?.shape.rect) cursor = getRectCornerCursor(focused.shape.rect, cursor_x, cursor_y) || cursor
-    else if (intersected?.isFocused && intersected.shape.polygon) cursor = 'move'
+    else if (intersected?.isFocused && intersected.shape.polygon) cursor = getPolygonPointCursor(intersected.shape.polygon.items, cursor_x, cursor_y) || 'move'
     else if (focused?.shape.polygon) cursor = getPolygonPointCursor(focused.shape.polygon.items, cursor_x, cursor_y) || cursor
 
     return cursor
@@ -260,8 +260,16 @@ export const XImageAnnotator = ({ model }: { model: ImageAnnotator }) => {
         }
         case 'select': {
           if (intersected) setActiveTag(intersected.tag)
+          if (intersected?.shape.polygon) polygonRef.current?.addAuxPoint(cursor_x, cursor_y, intersected.shape.polygon.items)
           polygonRef.current?.resetDragging()
-          setDrawnShapes(drawnShapes => drawnShapes.map(s => { s.isFocused = s === intersected; return s }))
+
+          setDrawnShapes(drawnShapes => drawnShapes.map(s => {
+            s.isFocused = s === intersected
+            if (s.isFocused && s.shape.polygon && polygonRef.current) {
+              s.shape.polygon.items = polygonRef.current.getPolygonPointsWithAux(s.shape.polygon.items)
+            }
+            return s
+          }))
           redrawExistingShapes()
           break
         }
@@ -334,7 +342,9 @@ export const XImageAnnotator = ({ model }: { model: ImageAnnotator }) => {
         tag,
         shape: {
           polygon: {
-            items: shape.polygon.items.map(i => ({ x: i.x / aspectRatio, y: i.y / aspectRatio }))
+            items: shape.polygon.items
+              .filter((i: DrawnPoint) => !i.isAux)
+              .map(i => ({ x: i.x / aspectRatio, y: i.y / aspectRatio }))
           }
         }
       }

diff --git a/ui/src/image_annotator_polygon.ts b/ui/src/image_annotator_polygon.ts
@@ -56,37 +56,51 @@ export class PolygonAnnotator {
     }
   }
 
-  drawLine = (x2: F, y2: F) => {
-    if (!this.ctx) return
-
-    this.ctx.lineTo(x2, y2)
-    this.ctx.stroke()
+  addAuxPoint = (cursor_x: F, cursor_y: F, items: DrawnPoint[]) => {
+    const clickedPoint = items.find(p => isIntersectingPoint(p, cursor_x, cursor_y))
+    if (clickedPoint?.isAux) clickedPoint.isAux = false
   }
 
-  drawPolygon = (points: DrawnPoint[], color: S, joinLastPoint = true, isFocused = false) => {
-    if (!points.length || !this.ctx) return
-    if (joinLastPoint && isFocused) {
-      points = points.reduce((prev, curr, idx) => {
-        if (!curr.isAux) prev.push(curr)
+  getPolygonPointsWithAux = (points: DrawnPoint[]) => {
+    const items = points
+      .filter(p => !p.isAux)
+      .reduce((prev, curr, idx, arr) => {
+        prev.push(curr)
 
-        if (idx !== points.length - 1 && !curr.isAux)
+        if (idx !== arr.length - 1) {
           prev.push({
-            x: (curr.x + points[idx + 1].x) / 2,
-            y: (curr.y + points[idx + 1].y) / 2,
+            x: (curr.x + arr[idx + 1].x) / 2,
+            y: (curr.y + arr[idx + 1].y) / 2,
             isAux: true
           })
+        }
 
         return prev
       }, [] as DrawnPoint[])
 
-      // Insert aux also between last and first point.
-      points.push({
-        x: (points[0].x + points.at(-1)!.x) / 2,
-        y: (points[0].y + points.at(-1)!.y) / 2,
+    // Insert aux also between last and first point.
+    const lastPoint = points.at(-1)?.isAux ? points.at(-2) : points.at(-1)
+    if (lastPoint) {
+      items.push({
+        x: (points[0].x + lastPoint.x) / 2,
+        y: (points[0].y + lastPoint.y) / 2,
         isAux: true
       })
     }
 
+    return items
+  }
+
+  drawLine = (x2: F, y2: F) => {
+    if (!this.ctx) return
+
+    this.ctx.lineTo(x2, y2)
+    this.ctx.stroke()
+  }
+
+  drawPolygon = (points: DrawnPoint[], color: S, joinLastPoint = true, isFocused = false) => {
+    if (!points.length || !this.ctx) return
+
     this.ctx.fillStyle = color
     this.ctx.strokeStyle = color
     this.ctx.beginPath()
@@ -122,7 +136,7 @@ export class PolygonAnnotator {
     path.arc(x, y, ARC_RADIUS, 0, 2 * Math.PI)
     this.ctx.lineWidth = 2
     this.ctx.strokeStyle = isAux ? '#5e5c5c' : '#000'
-    this.ctx.fillStyle = isAux ? '#e6e6e6' : '#FFF'
+    this.ctx.fillStyle = isAux ? '#b8b8b8' : '#FFF'
     this.ctx.fill(path)
     this.ctx.stroke(path)
   }
@@ -157,7 +171,11 @@ export
     const offset = 2 * ARC_RADIUS
     return cursor_x >= x - offset && cursor_x <= x + offset && cursor_y >= y - offset && cursor_y < y + offset
   },
-  getPolygonPointCursor = (items: ImageAnnotatorPoint[], cursor_x: F, cursor_y: F) => {
-    const isIntersecting = items.some(p => isIntersectingPoint(p, cursor_x, cursor_y))
-    return isIntersecting ? 'move' : ''
+  getPolygonPointCursor = (items: DrawnPoint[], cursor_x: F, cursor_y: F) => {
+    const intersectedPoint = items.find(p => isIntersectingPoint(p, cursor_x, cursor_y))
+    return intersectedPoint?.isAux
+      ? 'pointer'
+      : intersectedPoint
+        ? 'move'
+        : ''
   }