Add render_openposes_json API endpoint (#2325)

* Add render_openpose_json API endpoint

* nits

Author: huchenlei

annotator/openpose/__init__.py

@@ -62,23 +62,19 @@ def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, dr
     return canvas
 
 
-def decode_json_as_poses(json_string: str, normalize_coords: bool = False) -> Tuple[List[PoseResult], int, int]:
+def decode_json_as_poses(pose_json: dict) -> Tuple[List[PoseResult], int, int]:
     """ Decode the json_string complying with the openpose JSON output format
     to poses that controlnet recognizes.
     https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
 
     Args:
         json_string: The json string to decode.
-        normalize_coords: Whether to normalize coordinates of each keypoint by canvas height/width.
-                          `draw_pose` only accepts normalized keypoints. Set this param to True if
-                          the input coords are not normalized.
-    
+
     Returns:
         poses
         canvas_height
         canvas_width                      
     """
-    pose_json = json.loads(json_string)
     height = pose_json['canvas_height']
     width = pose_json['canvas_width']
 

scripts/api.py

@@ -1,8 +1,10 @@
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 from fastapi import FastAPI, Body
 from fastapi.exceptions import HTTPException
+from pydantic import BaseModel
+
 from PIL import Image
 
 import gradio as gr
@@ -13,6 +15,7 @@
 from scripts import external_code, global_state
 from scripts.processor import preprocessor_filters
 from scripts.logging import logger
+from annotator.openpose import draw_poses, decode_json_as_poses
 
 
 def encode_to_base64(image):
@@ -148,6 +151,32 @@ def accept(self, json_dict: dict) -> None:
 
         return res
 
+    class Person(BaseModel):
+        pose_keypoints_2d: List[float]
+        hand_right_keypoints_2d: Optional[List[float]]
+        hand_left_keypoints_2d: Optional[List[float]]
+        face_keypoints_2d: Optional[List[float]]
+
+    class PoseData(BaseModel):
+        people: List[Person]
+        canvas_width: int
+        canvas_height: int
+
+    @app.post("/controlnet/render_openpose_json")
+    async def render_openpose_json(
+        pose_data: List[PoseData] = Body([], title="Pose json files to render.")
+    ):
+        if not pose_data:
+            return {"info": "No pose data detected."}
+        else:
+            return {
+                "images": [
+                    encode_to_base64(draw_poses(*decode_json_as_poses(pose.dict())))
+                    for pose in pose_data
+                ],
+                "info": "Success",
+            }
+
 
 try:
     import modules.script_callbacks as script_callbacks

scripts/controlnet_ui/openpose_editor.py

@@ -1,5 +1,6 @@
 import base64
 import gradio as gr
+import json
 from typing import List, Dict, Any, Tuple
 
 from annotator.openpose import decode_json_as_poses, draw_poses
@@ -80,7 +81,7 @@ def register_callbacks(
     ):
         def render_pose(pose_url: str) -> Tuple[Dict, Dict]:
             json_string = parse_data_url(pose_url).decode('utf-8')
-            poses, height, weight = decode_json_as_poses(json_string)
+            poses, height, weight = decode_json_as_poses(json.loads(json_string))
             logger.info("Preview as input is enabled.")
             return (
                 # Generated image.

tests/web_api/pose.json

@@ -0,0 +1,194 @@
+{
+  "people": [
+    {
+      "pose_keypoints_2d": [
+        275.2506064884899,
+        196.32469357280343,
+        1,
+        303.3188016469506,
+        272.70982071889466,
+        1,
+        244.98447950024644,
+        292.09994638829477,
+        1,
+        236.38292745027104,
+        517.7037729015278,
+        1,
+        168.3984479500246,
+        418.0022744632577,
+        1,
+        412.90526047751445,
+        257.2121425016039,
+        1,
+        403.17894535813576,
+        510.14290732520925,
+        1,
+        294.43481869004336,
+        376.4781345848482,
+        1,
+        265.25747216047955,
+        562.5137576822758,
+        1,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        359.0078938961762,
+        562.0711206495608,
+        1,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        240.097671925037,
+        184.513914838073,
+        1,
+        308.33409148775263,
+        161.22089906296208,
+        1,
+        204.7558201874076,
+        213.05887067565308,
+        1,
+        366.61934701298674,
+        148.9278832878512,
+        1
+      ],
+      "hand_right_keypoints_2d": [
+        168.39790150130915,
+        418.0005271461072,
+        1,
+        181.79401055357368,
+        399.3767976307846,
+        1,
+        184.8627576873498,
+        384.75709227716266,
+        1,
+        198.118414869015,
+        381.90007483819153,
+        1,
+        215.15048903799024,
+        386.8527024571639,
+        1,
+        180.1325238303079,
+        346.88417988394843,
+        1,
+        178.10795487568174,
+        321.1018085790239,
+        1,
+        190.70710955474613,
+        320.5669160600145,
+        1,
+        203.3062827659017,
+        325.5456061326966,
+        1,
+        172.3536896669116,
+        350.7525276652412,
+        1,
+        170.000622912838,
+        325.0336533262108,
+        1,
+        185.70972426755964,
+        323.476117950832,
+        1,
+        208.50724912413568,
+        333.4635128502484,
+        1,
+        163.50256975319706,
+        356.1325737292637,
+        1,
+        162.59147123450128,
+        335.0197021116338,
+        1,
+        183.9828354600188,
+        328.4553078224726,
+        1,
+        201.57171021013423,
+        337.94383954551654,
+        1,
+        152.9805889462092,
+        357.94403689402753,
+        1,
+        167.09651929267878,
+        341.7437601311937,
+        1,
+        180.9402668216081,
+        337.10207327446744,
+        1,
+        194.60028340222678,
+        343.0449246874465,
+        1
+      ],
+      "hand_left_keypoints_2d": [
+        294.4393772120137,
+        376.476024395234,
+        1,
+        271.70933825161165,
+        384.48117305399165,
+        1,
+        257.2452829806548,
+        374.58948859472207,
+        1,
+        238.26122936397638,
+        375.2887100029166,
+        1,
+        219.89983184668415,
+        382.69322630254595,
+        1,
+        263.0323651487124,
+        320.1279349241104,
+        1,
+        246.94602107917282,
+        309.8099960810156,
+        1,
+        233.73717716804694,
+        314.1485136789638,
+        1,
+        224.27755744411303,
+        322.7892154545116,
+        1,
+        264.97558037166135,
+        334.6319791090978,
+        1,
+        254.35598193615226,
+        315.5629746257517,
+        1,
+        238.25810853722876,
+        321.2812182403252,
+        1,
+        223.57727818251382,
+        328.39525394113423,
+        1,
+        278.15831452661644,
+        337.1533682086847,
+        1,
+        265.12624946042416,
+        323.3619430418993,
+        1,
+        250.5919197031302,
+        325.30525908324694,
+        1,
+        235.2500911122877,
+        332.6721359855453,
+        1,
+        285.9427695830851,
+        341.50671458478496,
+        1,
+        274.50497773130155,
+        333.1376809270594,
+        1,
+        261.49768784257105,
+        328.7012203257942,
+        1,
+        248.90495501067193,
+        332.0535195828255,
+        1
+      ]
+    }
+  ],
+  "canvas_width": 512,
+  "canvas_height": 512
+}

tests/web_api/render_openpose_json.py

@@ -0,0 +1,42 @@
+import requests
+import unittest
+import importlib
+import json
+from pathlib import Path
+
+utils = importlib.import_module("extensions.sd-webui-controlnet.tests.utils", "utils")
+
+
+def render(poses):
+    return requests.post(
+        utils.BASE_URL + "/controlnet/render_openpose_json", json=poses
+    ).json()
+
+
+with open(Path(__file__).parent / "pose.json", "r") as f:
+    pose = json.load(f)
+
+
+class TestDetectEndpointWorking(unittest.TestCase):
+    def test_render_single(self):
+        res = render([pose])
+        self.assertEqual(res["info"], "Success")
+        self.assertEqual(len(res["images"]), 1)
+
+    def test_render_multiple(self):
+        res = render([pose, pose])
+        self.assertEqual(res["info"], "Success")
+        self.assertEqual(len(res["images"]), 2)
+
+    def test_render_no_pose(self):
+        res = render([])
+        self.assertNotEqual(res["info"], "Success")
+
+    def test_render_invalid_pose(self):
+        res = render([{"foo": 10, "bar": 100}])
+        self.assertNotIn("info", res)
+        self.assertNotIn("images", res)
+
+
+if __name__ == "__main__":
+    unittest.main()