Spaces:
Running
Running
| 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"): | |
| def __init__(self, model_path="./checkpoints/densepose_", device="cpu"): | |
| 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 | |
| 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. | |
| """ | |
| # random tmp path with timestamp | |
| 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") # predictor expects BGR image. | |
| # resize | |
| 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) | |
| # remove image_path and output_path | |
| os.remove(image_path) | |
| os.remove(output_path) | |
| return dense_gray | |
| if __name__ == '__main__': | |
| pass | |