Changes for statless GPU running torch.cuda (#5)

- Changes for statless GPU running torch.cuda (2d417ae15dfd712ecb8efe0cb1c8bc1d1be14df0)


Co-authored-by: rao <Balaji23@users.noreply.huggingface.co>

preprocess/openpose/run_openpose.py

@@ -1,84 +1,84 @@
-import pdb
-
-# import config
-from pathlib import Path
-import sys
-
-PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
-sys.path.insert(0, str(PROJECT_ROOT))
-import os
-
-import cv2
-import einops
-import numpy as np
-import random
-import time
-import json
-
-# from pytorch_lightning import seed_everything
-from preprocess.openpose.annotator.util import resize_image, HWC3
-from preprocess.openpose.annotator.openpose import OpenposeDetector
-
-import argparse
-from PIL import Image
-import torch
-import pdb
-
-# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
-
-class OpenPose:
-    def __init__(self, gpu_id: int):
-        self.gpu_id = gpu_id
-        torch.cuda.set_device(gpu_id)
-        self.preprocessor = OpenposeDetector()
-
-    def __call__(self, input_image, resolution=384):
-        torch.cuda.set_device(self.gpu_id)
-        if isinstance(input_image, Image.Image):
-            input_image = np.asarray(input_image)
-        elif type(input_image) == str:
-            input_image = np.asarray(Image.open(input_image))
-        else:
-            raise ValueError
-        with torch.no_grad():
-            input_image = HWC3(input_image)
-            input_image = resize_image(input_image, resolution)
-            H, W, C = input_image.shape
-            assert (H == 512 and W == 384), 'Incorrect input image shape'
-            pose, detected_map = self.preprocessor(input_image, hand_and_face=False)
-
-            candidate = pose['bodies']['candidate']
-            subset = pose['bodies']['subset'][0][:18]
-            for i in range(18):
-                if subset[i] == -1:
-                    candidate.insert(i, [0, 0])
-                    for j in range(i, 18):
-                        if(subset[j]) != -1:
-                            subset[j] += 1
-                elif subset[i] != i:
-                    candidate.pop(i)
-                    for j in range(i, 18):
-                        if(subset[j]) != -1:
-                            subset[j] -= 1
-
-            candidate = candidate[:18]
-
-            for i in range(18):
-                candidate[i][0] *= 384
-                candidate[i][1] *= 512
-
-            keypoints = {"pose_keypoints_2d": candidate}
-            # with open("/home/aigc/ProjectVTON/OpenPose/keypoints/keypoints.json", "w") as f:
-            #     json.dump(keypoints, f)
-            #
-            # # print(candidate)
-            # output_image = cv2.resize(cv2.cvtColor(detected_map, cv2.COLOR_BGR2RGB), (768, 1024))
-            # cv2.imwrite('/home/aigc/ProjectVTON/OpenPose/keypoints/out_pose.jpg', output_image)
-
-        return keypoints
-
-
-if __name__ == '__main__':
-
-    model = OpenPose()
-    model('./images/bad_model.jpg')
+import pdb
+
+# import config
+from pathlib import Path
+import sys
+
+PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+import os
+
+import cv2
+import einops
+import numpy as np
+import random
+import time
+import json
+
+# from pytorch_lightning import seed_everything
+from preprocess.openpose.annotator.util import resize_image, HWC3
+from preprocess.openpose.annotator.openpose import OpenposeDetector
+
+import argparse
+from PIL import Image
+import torch
+import pdb
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
+
+class OpenPose:
+    def __init__(self, gpu_id: int):
+        # self.gpu_id = gpu_id
+        # torch.cuda.set_device(gpu_id)
+        self.preprocessor = OpenposeDetector()
+
+    def __call__(self, input_image, resolution=384):
+        # torch.cuda.set_device(self.gpu_id)
+        if isinstance(input_image, Image.Image):
+            input_image = np.asarray(input_image)
+        elif type(input_image) == str:
+            input_image = np.asarray(Image.open(input_image))
+        else:
+            raise ValueError
+        with torch.no_grad():
+            input_image = HWC3(input_image)
+            input_image = resize_image(input_image, resolution)
+            H, W, C = input_image.shape
+            assert (H == 512 and W == 384), 'Incorrect input image shape'
+            pose, detected_map = self.preprocessor(input_image, hand_and_face=False)
+
+            candidate = pose['bodies']['candidate']
+            subset = pose['bodies']['subset'][0][:18]
+            for i in range(18):
+                if subset[i] == -1:
+                    candidate.insert(i, [0, 0])
+                    for j in range(i, 18):
+                        if(subset[j]) != -1:
+                            subset[j] += 1
+                elif subset[i] != i:
+                    candidate.pop(i)
+                    for j in range(i, 18):
+                        if(subset[j]) != -1:
+                            subset[j] -= 1
+
+            candidate = candidate[:18]
+
+            for i in range(18):
+                candidate[i][0] *= 384
+                candidate[i][1] *= 512
+
+            keypoints = {"pose_keypoints_2d": candidate}
+            # with open("/home/aigc/ProjectVTON/OpenPose/keypoints/keypoints.json", "w") as f:
+            #     json.dump(keypoints, f)
+            #
+            # # print(candidate)
+            # output_image = cv2.resize(cv2.cvtColor(detected_map, cv2.COLOR_BGR2RGB), (768, 1024))
+            # cv2.imwrite('/home/aigc/ProjectVTON/OpenPose/keypoints/out_pose.jpg', output_image)
+
+        return keypoints
+
+
+if __name__ == '__main__':
+
+    model = OpenPose()
+    model('./images/bad_model.jpg')