feat: Merge pull request #93 from UniversalDataTool/pose-estimation

Pose / Keypoint Annotation Interface

Author: seveibar

package.json

@@ -76,7 +76,7 @@
     "babel-loader": "^8.0.5",
     "babel-preset-react-app": "^7.0.0",
     "gh-pages": "^2.0.1",
-    "prettier": "^2.0.5",
+    "prettier": "^2.1.1",
     "raw.macro": "^0.3.0",
     "react-github-btn": "^1.1.1",
     "react-scripts": "^3.4.1"

src/Annotator/index.js

@@ -10,6 +10,7 @@ import type {
   MainLayoutState,
   Action,
 } from "../MainLayout/types"
+import type { KeypointsDefinition } from "../ImageCanvas/region-tools"
 import SettingsProvider from "../SettingsProvider"
 
 import combineReducers from "./reducers/combine-reducers.js"
@@ -41,6 +42,7 @@ type Props = {
   videoSrc?: string,
   keyframes?: Object,
   videoName?: string,
+  keypointDefinitions: KeypointsDefinition,
   fullImageSegmentationMode?: boolean,
   autoSegmentationOptions?:
     | {| type: "simple" |}
@@ -77,6 +79,7 @@ export const Annotator = ({
   onExit,
   onNextImage,
   onPrevImage,
+  keypointDefinitions,
   autoSegmentationOptions = { type: "autoseg" },
 }: Props) => {
   if (typeof selectedImage === "string") {
@@ -112,6 +115,7 @@ export const Annotator = ({
       enabledTools,
       history: [],
       videoName,
+      keypointDefinitions,
       ...(annotationType === "image"
         ? {
             selectedImage,

src/Annotator/index.story.js

@@ -441,7 +441,7 @@ storiesOf("Annotator", module)
         videoSrc="https://s3.amazonaws.com/asset.workaround.online/SampleVideo_1280x720_1mb.mp4"
         videoTime={1000}
         keyframes={{
-          "0": {
+          0: {
             regions: [
               {
                 type: "point",
@@ -465,7 +465,7 @@ storiesOf("Annotator", module)
               },
             ],
           },
-          "3333": {
+          3333: {
             regions: [
               {
                 type: "point",

src/Annotator/poses.story.js

@@ -0,0 +1,76 @@
+// @flow
+
+import React, { useState } from "react"
+
+import { storiesOf } from "@storybook/react"
+import { action as actionAddon } from "@storybook/addon-actions"
+import dancingManImage from "../ImageCanvas/dancing-man.story.jpg"
+import Annotator from "./"
+
+const middlewares = [
+  (store) => (next) => (action) => {
+    actionAddon(action.type)(action)
+    return next(action)
+  },
+]
+
+storiesOf("Annotator (Poses)", module).add("Basic", () => (
+  <Annotator
+    onExit={actionAddon("onExit")}
+    middlewares={middlewares}
+    labelImages
+    enabledTools={["create-keypoints"]}
+    keypointDefinitions={{
+      human: {
+        connections: [
+          ["head", "sternum"],
+          ["sternum", "leftElbow"],
+          ["sternum", "rightElbow"],
+        ],
+        landmarks: {
+          head: {
+            label: "Head",
+            color: "#f00",
+            defaultPosition: [0, -0.05],
+          },
+          sternum: {
+            label: "Torso",
+            color: "#0f0",
+            defaultPosition: [0, 0],
+          },
+          leftElbow: {
+            label: "Left Elbow",
+            color: "#00f",
+            defaultPosition: [-0.05, 0],
+          },
+          rightElbow: {
+            label: "Right Elbow",
+            color: "#00f",
+            defaultPosition: [0.05, 0],
+          },
+        },
+      },
+    }}
+    images={[
+      {
+        src: dancingManImage,
+        name: "Dancing Man",
+        regions: [
+          {
+            type: "keypoints",
+            id: "keypoints1",
+            keypointsDefinitionId: "human",
+            highlighted: true,
+            points: {
+              head: { x: 0.54, y: 0.2 },
+              sternum: { x: 0.57, y: 0.3 },
+              leftElbow: { x: 0.4, y: 0.39 },
+              rightElbow: { x: 0.7, y: 0.32 },
+            },
+            visible: true,
+          },
+        ],
+      },
+    ]}
+  />
+))

src/Annotator/reducers/general-reducer.js

@@ -10,6 +10,7 @@ import colors from "../../colors"
 import fixTwisted from "./fix-twisted"
 import convertExpandingLineToPolygon from "./convert-expanding-line-to-polygon"
 import clamp from "clamp"
+import getLandmarksWithTransform from "../../utils/get-landmarks-with-transform"
 
 const getRandomId = () => Math.random().toString().split(".")[1]
 
@@ -206,6 +207,16 @@ export default (state: MainLayoutState, action: Action) => {
         pointIndex,
       })
     }
+    case "BEGIN_MOVE_KEYPOINT": {
+      const { region, keypointId } = action
+      state = closeEditors(state)
+      state = saveToHistory(state, "Move Keypoint")
+      return setIn(state, ["mode"], {
+        mode: "MOVE_KEYPOINT",
+        regionId: region.id,
+        keypointId,
+      })
+    }
     case "ADD_POLYGON_POINT": {
       const { polygon, point, pointIndex } = action
       const regionIndex = getRegionIndex(polygon)
@@ -240,6 +251,22 @@ export default (state: MainLayoutState, action: Action) => {
             [x, y]
           )
         }
+        case "MOVE_KEYPOINT": {
+          const { keypointId, regionId } = state.mode
+          const [region, regionIndex] = getRegion(regionId)
+          if (regionIndex === null) return state
+          return setIn(
+            state,
+            [
+              ...pathToActiveImage,
+              "regions",
+              regionIndex,
+              "points",
+              keypointId,
+            ],
+            { ...(region: any).points[keypointId], x, y }
+          )
+        }
         case "MOVE_REGION": {
           const { regionId } = state.mode
           if (regionId === "$$allowed_area") {
@@ -312,6 +339,20 @@ export default (state: MainLayoutState, action: Action) => {
             h: dh,
           })
         }
+        case "RESIZE_KEYPOINTS": {
+          const { regionId, landmarks, centerX, centerY } = state.mode
+          const distFromCenter = Math.sqrt(
+            (centerX - x) ** 2 + (centerY - y) ** 2
+          )
+          const scale = distFromCenter / 0.15
+          return modifyRegion(regionId, {
+            points: getLandmarksWithTransform({
+              landmarks,
+              center: { x: centerX, y: centerY },
+              scale,
+            }),
+          })
+        }
         case "DRAW_POLYGON": {
           const { regionId } = state.mode
           const [region, regionIndex] = getRegion(regionId)
@@ -541,6 +582,34 @@ export default (state: MainLayoutState, action: Action) => {
           })
           break
         }
+        case "create-keypoints": {
+          state = saveToHistory(state, "Create Keypoints")
+          const [
+            [keypointsDefinitionId, { landmarks, connections }],
+          ] = (Object.entries(state.keypointDefinitions): any)
+
+          newRegion = {
+            type: "keypoints",
+            keypointsDefinitionId,
+            points: getLandmarksWithTransform({
+              landmarks,
+              center: { x, y },
+              scale: 1,
+            }),
+            highlighted: true,
+            editingLabels: false,
+            id: getRandomId(),
+          }
+          state = setIn(state, ["mode"], {
+            mode: "RESIZE_KEYPOINTS",
+            landmarks,
+            centerX: x,
+            centerY: y,
+            regionId: newRegion.id,
+            isNew: true,
+          })
+          break
+        }
         default:
           break
       }
@@ -581,9 +650,13 @@ export default (state: MainLayoutState, action: Action) => {
           }
         }
         case "MOVE_REGION":
+        case "RESIZE_KEYPOINTS":
         case "MOVE_POLYGON_POINT": {
           return { ...state, mode: null }
         }
+        case "MOVE_KEYPOINT": {
+          return { ...state, mode: null }
+        }
         case "CREATE_POINT_LINE": {
           return state
         }

src/Annotator/video.story.js

@@ -11,7 +11,7 @@ storiesOf("Annotator(video)", module).add("Video Annotator", () => {
     regionClsList: ["valid", "invalid"],
     enabledTools: ["select", "create-box", "create-polygon", "create-point"],
     keyframes: {
-      "0": {
+      0: {
         regions: [
           {
             id: "910517662556203",
@@ -25,7 +25,7 @@ storiesOf("Annotator(video)", module).add("Video Annotator", () => {
           },
         ],
       },
-      "2656": {
+      2656: {
         regions: [
           {
             id: "910517662556203",

src/ImageCanvas/dancing-man.story.jpg

@@ -10,7 +10,14 @@ import React, {
 import type { Node } from "react"
 import { Matrix } from "transformation-matrix-js"
 import Crosshairs from "../Crosshairs"
-import type { Region, Point, Polygon, Box } from "./region-tools.js"
+import type {
+  Region,
+  Point,
+  Polygon,
+  Box,
+  Keypoints,
+  KeypointsDefinition,
+} from "./region-tools.js"
 import { makeStyles } from "@material-ui/core/styles"
 import styles from "./styles"
 import PreventScrollToParents from "../PreventScrollToParents"
@@ -36,6 +43,7 @@ type Props = {
   imageSrc?: string,
   videoSrc?: string,
   videoTime?: number,
+  keypointDefinitions?: KeypointDefinitions,
   onMouseMove?: ({ x: number, y: number }) => any,
   onMouseDown?: ({ x: number, y: number }) => any,
   onMouseUp?: ({ x: number, y: number }) => any,
@@ -65,6 +73,7 @@ type Props = {
   onDeleteRegion: (Region) => any,
   onBeginBoxTransform: (Box, [number, number]) => any,
   onBeginMovePolygonPoint: (Polygon, index: number) => any,
+  onBeginMoveKeypoint: (Keypoints, index: number) => any,
   onAddPolygonPoint: (Polygon, point: [number, number], index: number) => any,
   onSelectRegion: (Region) => any,
   onBeginMovePoint: (Point) => any,
@@ -120,6 +129,7 @@ export const ImageCanvas = ({
   onBeginBoxTransform,
   onBeginMovePolygonPoint,
   onAddPolygonPoint,
+  onBeginMoveKeypoint,
   onSelectRegion,
   onBeginMovePoint,
   onDeleteRegion,
@@ -128,6 +138,7 @@ export const ImageCanvas = ({
   onRegionClassAdded,
   zoomOnAllowedArea = true,
   modifyingAllowedArea = false,
+  keypointDefinitions,
 }: Props) => {
   const classes = useStyles()
 
@@ -356,6 +367,7 @@ export const ImageCanvas = ({
           mat={mat}
           onBeginBoxTransform={onBeginBoxTransform}
           onBeginMovePolygonPoint={onBeginMovePolygonPoint}
+          onBeginMoveKeypoint={onBeginMoveKeypoint}
           onAddPolygonPoint={onAddPolygonPoint}
           showHighlightBox={showHighlightBox}
         />
@@ -440,6 +452,7 @@ export const ImageCanvas = ({
           />
           <RegionShapes
             mat={mat}
+            keypointDefinitions={keypointDefinitions}
             imagePosition={imagePosition}
             regions={regions}
             fullSegmentationMode={fullImageSegmentationMode}