From 446c28671ba8a2b4e33813ff781d9ba144851cbc Mon Sep 17 00:00:00 2001
From: Hyunmin-H <handy389@naver.com>
Date: Wed, 16 Aug 2023 12:50:53 +0000
Subject: [PATCH] =?UTF-8?q?[Feat]=20Pose=20estimation=20I/O=20=ED=98=95?=
 =?UTF-8?q?=EC=8B=9D=20=EB=B3=80=EA=B2=BD=20#33?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- openpose의 Input 형식을 image byte로, output 형식을 dictionary
형식으로 변경

related to : #31
---
 model/pytorch_openpose/extract_keypoint.py | 43 ++++++++++++++++++++--
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/model/pytorch_openpose/extract_keypoint.py b/model/pytorch_openpose/extract_keypoint.py
index c7dd26c..4ed5d4a 100644
--- a/model/pytorch_openpose/extract_keypoint.py
+++ b/model/pytorch_openpose/extract_keypoint.py
@@ -12,9 +12,6 @@
 
 import json
 
-# TEST_IMAGE = '/opt/ml/dataset/ganddddi/sw2_0.jpg'
-# OUTPUT_JSON = '/opt/ml/pytorch-openpose/output/sw2_0_hi.json'
-
 # target_image 512, 384
 def main_openpose(target_buffer_dir, output_buffer_dir):
     
@@ -52,3 +49,43 @@ def main_openpose(target_buffer_dir, output_buffer_dir):
     with open(f'{output_buffer_dir}/{img_name[:-4]}.json', 'w') as f:
         json.dump(json_dict, f, indent=4)
 
+def main_openpose_fromImageByte(target_bytes):
+    
+    img_name = 'target.jpg'
+    
+    # Body, Hand model load
+    body_estimation = Body('/opt/ml/checkpoints/body_pose_model.pth')
+    # hand_estimation = Hand('model/hand_pose_model.pth')
+
+
+    # image read
+    # test_image = test_image
+    
+    # oriImg = cv2.imread(test_image)  # B,G,R order
+    nparr = np.frombuffer(target_bytes, np.uint8)
+    oriImg = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+    oriImg = cv2.resize(oriImg, (384, 512)) # resize
+
+    # body_estimation foreward
+    candidate, subset = body_estimation(oriImg)
+
+    while True:
+        if len(candidate) == 18:
+            break
+        elif len(candidate) > 18:
+            candidate = candidate[:-1]
+        elif len(candidate) < 18:
+            candidate = list(candidate)
+            candidate.append(np.array([-1.0, -1.0, 0.0, -1.0]))
+            candidate = np.array(candidate)
+
+    json_dict = {
+        'keypoints' : candidate.tolist()
+    }
+
+    # json 파일 저장
+    # with open(f'{output_buffer_dir}/{img_name[:-4]}.json', 'w') as f:
+    #     json.dump(json_dict, f, indent=4)
+    return json_dict
+