Spaces:
Running
Running
File size: 4,307 Bytes
b85284b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
from __future__ import annotations
import os
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: str):
self.config = self._load_config()
self.config['device'] = device
self._download_models()
self.model = SampleFromPoseModel(self.config)
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
token = os.getenv('HF_TOKEN')
path = huggingface_hub.hf_hub_download('hysts/Text2Human',
'orig/pretrained_models.zip',
use_auth_token=token)
model_dir.mkdir()
with zipfile.ZipFile(path) as f:
f.extractall(model_dir)
@staticmethod
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. - 1
data = torch.from_numpy(image).unsqueeze(1)
return data
@staticmethod
def process_mask(mask: torch.Tensor) -> torch.Tensor:
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
assert (seg_map != -1).all()
return seg_map
@staticmethod
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) -> None:
if pose_image is None:
return
data = self.preprocess_pose_image(pose_image)
self.model.feed_pose_data(data)
def generate_label_image(self, shape_text: str) -> np.ndarray:
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())
mask = colored_segm.copy()
seg_map = self.process_mask(mask)
self.model.segm = torch.from_numpy(seg_map).unsqueeze(0).unsqueeze(
0).to(self.model.device)
self.model.generate_quantized_segm()
return colored_segm
def generate_human(self, texture_text: str, sample_steps: int,
seed: int) -> np.ndarray:
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
|