|
|
| import glob |
| import os |
| from random import randint |
| import shutil |
| import time |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| from PIL import Image |
| from densepose import add_densepose_config |
| from densepose.vis.base import CompoundVisualizer |
| from densepose.vis.densepose_results import DensePoseResultsFineSegmentationVisualizer |
| from densepose.vis.extractor import create_extractor, CompoundExtractor |
| from detectron2.config import get_cfg |
| from detectron2.data.detection_utils import read_image |
| from detectron2.engine.defaults import DefaultPredictor |
|
|
|
|
| class DensePose: |
| """ |
| DensePose used in this project is from Detectron2 (https://github.com/facebookresearch/detectron2). |
| These codes are modified from https://github.com/facebookresearch/detectron2/tree/main/projects/DensePose. |
| The checkpoint is downloaded from https://github.com/facebookresearch/detectron2/blob/main/projects/DensePose/doc/DENSEPOSE_IUV.md#ModelZoo. |
| |
| We use the model R_50_FPN_s1x with id 165712039, but other models should also work. |
| The config file is downloaded from https://github.com/facebookresearch/detectron2/tree/main/projects/DensePose/configs. |
| Noted that the config file should match the model checkpoint and Base-DensePose-RCNN-FPN.yaml is also needed. |
| """ |
|
|
| def __init__(self, model_path="./checkpoints/densepose_", device="cuda"): |
| self.device = device |
| self.config_path = os.path.join(model_path, 'densepose_rcnn_R_50_FPN_s1x.yaml') |
| self.model_path = os.path.join(model_path, 'model_final_162be9.pkl') |
| self.visualizations = ["dp_segm"] |
| self.VISUALIZERS = {"dp_segm": DensePoseResultsFineSegmentationVisualizer} |
| self.min_score = 0.8 |
|
|
| self.cfg = self.setup_config() |
| self.predictor = DefaultPredictor(self.cfg) |
| self.predictor.model.to(self.device) |
|
|
| def setup_config(self): |
| opts = ["MODEL.ROI_HEADS.SCORE_THRESH_TEST", str(self.min_score)] |
| cfg = get_cfg() |
| add_densepose_config(cfg) |
| cfg.merge_from_file(self.config_path) |
| cfg.merge_from_list(opts) |
| cfg.MODEL.WEIGHTS = self.model_path |
| cfg.freeze() |
| return cfg |
|
|
| @staticmethod |
| def _get_input_file_list(input_spec: str): |
| if os.path.isdir(input_spec): |
| file_list = [os.path.join(input_spec, fname) for fname in os.listdir(input_spec) |
| if os.path.isfile(os.path.join(input_spec, fname))] |
| elif os.path.isfile(input_spec): |
| file_list = [input_spec] |
| else: |
| file_list = glob.glob(input_spec) |
| return file_list |
|
|
| def create_context(self, cfg, output_path): |
| vis_specs = self.visualizations |
| visualizers = [] |
| extractors = [] |
| for vis_spec in vis_specs: |
| texture_atlas = texture_atlases_dict = None |
| vis = self.VISUALIZERS[vis_spec]( |
| cfg=cfg, |
| texture_atlas=texture_atlas, |
| texture_atlases_dict=texture_atlases_dict, |
| alpha=1.0 |
| ) |
| visualizers.append(vis) |
| extractor = create_extractor(vis) |
| extractors.append(extractor) |
| visualizer = CompoundVisualizer(visualizers) |
| extractor = CompoundExtractor(extractors) |
| context = { |
| "extractor": extractor, |
| "visualizer": visualizer, |
| "out_fname": output_path, |
| "entry_idx": 0, |
| } |
| return context |
|
|
| def execute_on_outputs(self, context, entry, outputs): |
| extractor = context["extractor"] |
|
|
| data = extractor(outputs) |
|
|
| H, W, _ = entry["image"].shape |
| result = np.zeros((H, W), dtype=np.uint8) |
|
|
| data, box = data[0] |
| x, y, w, h = [int(_) for _ in box[0].cpu().numpy()] |
| i_array = data[0].labels[None].cpu().numpy()[0] |
| result[y:y + h, x:x + w] = i_array |
| result = Image.fromarray(result) |
| result.save(context["out_fname"]) |
|
|
| def __call__(self, image_or_path, resize=512) -> Image.Image: |
| """ |
| :param image_or_path: Path of the input image. |
| :param resize: Resize the input image if its max size is larger than this value. |
| :return: Dense pose image. |
| """ |
| |
| tmp_path = f"./densepose_/tmp/" |
| if not os.path.exists(tmp_path): |
| os.makedirs(tmp_path) |
|
|
| image_path = os.path.join(tmp_path, f"{int(time.time())}-{self.device}-{randint(0, 100000)}.png") |
| if isinstance(image_or_path, str): |
| assert image_or_path.split(".")[-1] in ["jpg", "png"], "Only support jpg and png images." |
| shutil.copy(image_or_path, image_path) |
| elif isinstance(image_or_path, Image.Image): |
| image_or_path.save(image_path) |
| else: |
| shutil.rmtree(tmp_path) |
| raise TypeError("image_path must be str or PIL.Image.Image") |
|
|
| output_path = image_path.replace(".png", "_dense.png").replace(".jpg", "_dense.png") |
| w, h = Image.open(image_path).size |
|
|
| file_list = self._get_input_file_list(image_path) |
| assert len(file_list), "No input images found!" |
| context = self.create_context(self.cfg, output_path) |
| for file_name in file_list: |
| img = read_image(file_name, format="BGR") |
| |
| if (_ := max(img.shape)) > resize: |
| scale = resize / _ |
| img = cv2.resize(img, (int(img.shape[1] * scale), int(img.shape[0] * scale))) |
|
|
| with torch.no_grad(): |
| outputs = self.predictor(img)["instances"] |
| try: |
| self.execute_on_outputs(context, {"file_name": file_name, "image": img}, outputs) |
| except Exception as e: |
| null_gray = Image.new('L', (1, 1)) |
| null_gray.save(output_path) |
|
|
| dense_gray = Image.open(output_path).convert("L") |
| dense_gray = dense_gray.resize((w, h), Image.NEAREST) |
| |
| os.remove(image_path) |
| os.remove(output_path) |
|
|
|
|
| return dense_gray |
|
|
|
|
| if __name__ == '__main__': |
| pass |
|
|
