Spaces:
Running
Running
| from __future__ import annotations | |
| import pathlib | |
| import sys | |
| import zipfile | |
| import huggingface_hub | |
| import numpy as np | |
| import PIL.Image | |
| import torch | |
| sys.path.insert(0, "Text2Human") | |
| from models.sample_model import SampleFromPoseModel | |
| from utils.language_utils import generate_shape_attributes, generate_texture_attributes | |
| from utils.options import dict_to_nonedict, parse | |
| from utils.util import set_random_seed | |
| COLOR_LIST = [ | |
| (0, 0, 0), | |
| (255, 250, 250), | |
| (220, 220, 220), | |
| (250, 235, 215), | |
| (255, 250, 205), | |
| (211, 211, 211), | |
| (70, 130, 180), | |
| (127, 255, 212), | |
| (0, 100, 0), | |
| (50, 205, 50), | |
| (255, 255, 0), | |
| (245, 222, 179), | |
| (255, 140, 0), | |
| (255, 0, 0), | |
| (16, 78, 139), | |
| (144, 238, 144), | |
| (50, 205, 174), | |
| (50, 155, 250), | |
| (160, 140, 88), | |
| (213, 140, 88), | |
| (90, 140, 90), | |
| (185, 210, 205), | |
| (130, 165, 180), | |
| (225, 141, 151), | |
| ] | |
| class Model: | |
| def __init__(self): | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| self.config = self._load_config() | |
| self.config["device"] = device.type | |
| self._download_models() | |
| self.model = SampleFromPoseModel(self.config) | |
| self.model.batch_size = 1 | |
| def _load_config(self) -> dict: | |
| path = "Text2Human/configs/sample_from_pose.yml" | |
| config = parse(path, is_train=False) | |
| config = dict_to_nonedict(config) | |
| return config | |
| def _download_models(self) -> None: | |
| model_dir = pathlib.Path("pretrained_models") | |
| if model_dir.exists(): | |
| return | |
| path = huggingface_hub.hf_hub_download("yumingj/Text2Human_SSHQ", "pretrained_models.zip") | |
| model_dir.mkdir() | |
| with zipfile.ZipFile(path) as f: | |
| f.extractall(model_dir) | |
| def preprocess_pose_image(image: PIL.Image.Image) -> torch.Tensor: | |
| image = ( | |
| np.array(image.resize(size=(256, 512), resample=PIL.Image.Resampling.LANCZOS))[:, :, 2:] | |
| .transpose(2, 0, 1) | |
| .astype(np.float32) | |
| ) | |
| image = image / 12.0 - 1 | |
| data = torch.from_numpy(image).unsqueeze(1) | |
| return data | |
| def process_mask(mask: np.ndarray) -> np.ndarray: | |
| if mask.shape != (512, 256, 3): | |
| return None | |
| seg_map = np.full(mask.shape[:-1], -1) | |
| for index, color in enumerate(COLOR_LIST): | |
| seg_map[np.sum(mask == color, axis=2) == 3] = index | |
| if not (seg_map != -1).all(): | |
| return None | |
| return seg_map | |
| def postprocess(result: torch.Tensor) -> np.ndarray: | |
| result = result.permute(0, 2, 3, 1) | |
| result = result.detach().cpu().numpy() | |
| result = result * 255 | |
| result = np.asarray(result[0, :, :, :], dtype=np.uint8) | |
| return result | |
| def process_pose_image(self, pose_image: PIL.Image.Image) -> torch.Tensor: | |
| if pose_image is None: | |
| return | |
| data = self.preprocess_pose_image(pose_image) | |
| self.model.feed_pose_data(data) | |
| return data | |
| def generate_label_image(self, pose_data: torch.Tensor, shape_text: str) -> np.ndarray: | |
| if pose_data is None: | |
| return | |
| self.model.feed_pose_data(pose_data) | |
| shape_attributes = generate_shape_attributes(shape_text) | |
| shape_attributes = torch.LongTensor(shape_attributes).unsqueeze(0) | |
| self.model.feed_shape_attributes(shape_attributes) | |
| self.model.generate_parsing_map() | |
| self.model.generate_quantized_segm() | |
| colored_segm = self.model.palette_result(self.model.segm[0].cpu()) | |
| return colored_segm | |
| def generate_human(self, label_image: np.ndarray, texture_text: str, sample_steps: int, seed: int) -> np.ndarray: | |
| if label_image is None: | |
| return | |
| mask = label_image.copy() | |
| seg_map = self.process_mask(mask) | |
| if seg_map is None: | |
| return | |
| self.model.segm = torch.from_numpy(seg_map).unsqueeze(0).unsqueeze(0).to(self.model.device) | |
| self.model.generate_quantized_segm() | |
| set_random_seed(seed) | |
| texture_attributes = generate_texture_attributes(texture_text) | |
| texture_attributes = torch.LongTensor(texture_attributes) | |
| self.model.feed_texture_attributes(texture_attributes) | |
| self.model.generate_texture_map() | |
| self.model.sample_steps = sample_steps | |
| out = self.model.sample_and_refine() | |
| res = self.postprocess(out) | |
| return res | |