Add mmpose

model.py

@@ -41,6 +41,73 @@ sys.path.append('T2I-Adapter')
 config_path =  'https://github.com/TencentARC/T2I-Adapter/raw/main/configs/stable-diffusion/'
 model_path = 'https://github.com/TencentARC/T2I-Adapter/raw/main/models/'
 
+
+def imshow_keypoints(img,
+                     pose_result,
+                     skeleton=None,
+                     kpt_score_thr=0.1,
+                     pose_kpt_color=None,
+                     pose_link_color=None,
+                     radius=4,
+                     thickness=1):
+    """Draw keypoints and links on an image.
+    Args:
+            img (ndarry): The image to draw poses on.
+            pose_result (list[kpts]): The poses to draw. Each element kpts is
+                a set of K keypoints as an Kx3 numpy.ndarray, where each
+                keypoint is represented as x, y, score.
+            kpt_score_thr (float, optional): Minimum score of keypoints
+                to be shown. Default: 0.3.
+            pose_kpt_color (np.array[Nx3]`): Color of N keypoints. If None,
+                the keypoint will not be drawn.
+            pose_link_color (np.array[Mx3]): Color of M links. If None, the
+                links will not be drawn.
+            thickness (int): Thickness of lines.
+    """
+
+    img_h, img_w, _ = img.shape
+    img = np.zeros(img.shape)
+
+    for idx, kpts in enumerate(pose_result):
+        if idx > 1:
+            continue
+        kpts = kpts['keypoints']
+        # print(kpts)
+        kpts = np.array(kpts, copy=False)
+
+        # draw each point on image
+        if pose_kpt_color is not None:
+            assert len(pose_kpt_color) == len(kpts)
+
+            for kid, kpt in enumerate(kpts):
+                x_coord, y_coord, kpt_score = int(kpt[0]), int(kpt[1]), kpt[2]
+
+                if kpt_score < kpt_score_thr or pose_kpt_color[kid] is None:
+                    # skip the point that should not be drawn
+                    continue
+
+                color = tuple(int(c) for c in pose_kpt_color[kid])
+                cv2.circle(img, (int(x_coord), int(y_coord)), radius, color, -1)
+
+        # draw links
+        if skeleton is not None and pose_link_color is not None:
+            assert len(pose_link_color) == len(skeleton)
+
+            for sk_id, sk in enumerate(skeleton):
+                pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1]))
+                pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1]))
+
+                if (pos1[0] <= 0 or pos1[0] >= img_w or pos1[1] <= 0 or pos1[1] >= img_h or pos2[0] <= 0
+                        or pos2[0] >= img_w or pos2[1] <= 0 or pos2[1] >= img_h or kpts[sk[0], 2] < kpt_score_thr
+                        or kpts[sk[1], 2] < kpt_score_thr or pose_link_color[sk_id] is None):
+                    # skip the link that should not be drawn
+                    continue
+                color = tuple(int(c) for c in pose_link_color[sk_id])
+                cv2.line(img, pos1, pos2, color, thickness=thickness)
+
+    return img
+
+
 def load_model_from_config(config, ckpt, verbose=False):
     print(f"Loading model from {ckpt}")
     pl_sd = torch.load(ckpt, map_location="cpu")
@@ -71,10 +138,36 @@ class Model:
         self.device = torch.device(
             'cuda:0' if torch.cuda.is_available() else 'cpu')    
         self.model_dir = pathlib.Path(model_dir)
-
+        self.download_pose_models()
         self.download_models()
 
 
+    def download_pose_models(self) -> None:
+        ## mmpose
+        device = "cuda"
+        det_config_file = model_path+"faster_rcnn_r50_fpn_coco.py"
+        subprocess.run(shlex.split(f'wget {det_config_file} -O models/faster_rcnn_r50_fpn_coco.py'))
+        det_config = 'models/faster_rcnn_r50_fpn_coco.py' 
+        
+        det_checkpoint_file = "https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth"
+        subprocess.run(shlex.split(f'wget {det_checkpoint_file} -O models/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'))
+        det_checkpoint = 'models/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
+
+        pose_config_file = model_path+"rnet_w48_coco_256x192.py"
+        subprocess.run(shlex.split(f'wget {pose_config_file} -O models/rnet_w48_coco_256x192.py'))
+        pose_config = 'models/hrnet_w48_coco_256x192.py'
+
+        pose_checkpoint_file = "https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth"
+        subprocess.run(shlex.split(f'wget {pose_checkpoint_file} -O models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'))
+        pose_checkpoint = 'models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'
+        
+        det_cat_id = 1
+        bbox_thr = 0.2
+        ## detector
+        det_config_mmcv = mmcv.Config.fromfile(det_config)
+        self.det_model = init_detector(det_config_mmcv, det_checkpoint, device=device)
+        pose_config_mmcv = mmcv.Config.fromfile(pose_config)
+        self.pose_model = init_pose_model(pose_config_mmcv, pose_checkpoint, device=device)
 
     def download_models(self) -> None:
         self.model_dir.mkdir(exist_ok=True, parents=True)
@@ -206,16 +299,49 @@ class Model:
             seed_everything(42)
         
         im = cv2.resize(input_img,(512,512))
-        pose = img2tensor(im, bgr2rgb=True, float32=True)/255.
-        pose = pose.unsqueeze(0)
 
-        im_pose = tensor2img(pose)
+        image = im.copy()
+        im = img2tensor(im).unsqueeze(0)/255.
+        mmdet_results = inference_detector(det_model, image)
+        # keep the person class bounding boxes.
+        person_results = process_mmdet_results(mmdet_results, det_cat_id)
+
+        # optional
+        return_heatmap = False
+        dataset = pose_model.cfg.data['test']['type']
+
+        # e.g. use ('backbone', ) to return backbone feature
+        output_layer_names = None
+        pose_results, returned_outputs = inference_top_down_pose_model(
+            pose_model,
+            image,
+            person_results,
+            bbox_thr=bbox_thr,
+            format='xyxy',
+            dataset=dataset,
+            dataset_info=None,
+            return_heatmap=return_heatmap,
+            outputs=output_layer_names)
+
+        # show the results
+        im_pose = imshow_keypoints(
+            image,
+            pose_results,
+            skeleton=skeleton,
+            pose_kpt_color=pose_kpt_color,
+            pose_link_color=pose_link_color,
+            radius=2,
+            thickness=2)
+
+        im_pose = cv2.resize(im_pose,(512,512))
     
         c = model.get_learned_conditioning([prompt])
         nc = model.get_learned_conditioning([neg_prompt])
         
         with torch.no_grad():
             # extract condition features
+            pose = img2tensor(im_pose, bgr2rgb=True, float32=True)/255.
+            pose = pose.unsqueeze(0)
             features_adapter = self.model_ad_pose(pose.to(device))
     
         shape = [4, 64, 64]