Spaces:
Running
Running
File size: 6,960 Bytes
1080e5b 77c9ee7 1080e5b 77c9ee7 1080e5b 77c9ee7 1080e5b 77c9ee7 1080e5b |
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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 |
import os
import sys
import numpy as np
import torch
import yaml
import glob
import argparse
from omegaconf import OmegaConf
from pathlib import Path
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
sys.path.insert(0, str(Path(__file__).resolve().parent / "lama"))
from saicinpainting.evaluation.utils import move_to_device
from saicinpainting.training.trainers import load_checkpoint
from saicinpainting.evaluation.data import pad_tensor_to_modulo
from saicinpainting.evaluation.refinement import refine_predict
from utils import load_img_to_array, save_array_to_img
@torch.no_grad()
def inpaint_img_with_lama(
img: np.ndarray, mask: np.ndarray, config_p: str, ckpt_p: str, mod=8, device="cuda"
):
assert len(mask.shape) == 2
if np.max(mask) == 1:
mask = mask * 255
img = torch.from_numpy(img).float().div(255.0)
mask = torch.from_numpy(mask).float()
predict_config = OmegaConf.load(config_p)
predict_config.model.path = ckpt_p
# device = torch.device(predict_config.device)
device = torch.device(device)
train_config_path = os.path.join(predict_config.model.path, "config.yaml")
with open(train_config_path, "r") as f:
train_config = OmegaConf.create(yaml.safe_load(f))
train_config.training_model.predict_only = True
train_config.visualizer.kind = "noop"
checkpoint_path = os.path.join(
predict_config.model.path, "models", predict_config.model.checkpoint
)
model = load_checkpoint(
train_config, checkpoint_path, strict=False, map_location=device
)
model.freeze()
model.to(device)
batch = {}
batch["image"] = img.permute(2, 0, 1).unsqueeze(0)
batch["mask"] = mask[None, None]
unpad_to_size = [batch["image"].shape[2], batch["image"].shape[3]]
batch["image"] = pad_tensor_to_modulo(batch["image"], mod)
batch["mask"] = pad_tensor_to_modulo(batch["mask"], mod)
# batch = move_to_device(batch, device)
# batch["mask"] = (batch["mask"] > 0) * 1
# batch = model(batch)
# cur_res = batch[predict_config.out_key][0].permute(1, 2, 0)
# cur_res = cur_res.detach().cpu().numpy()
if predict_config.get("refine", False):
batch["unpad_to_size"] = [torch.tensor([size]) for size in unpad_to_size]
cur_res = refine_predict(batch, model, **predict_config.refiner)
cur_res = cur_res[0].permute(1, 2, 0).detach().cpu().numpy()
else:
batch = move_to_device(batch, device)
batch["mask"] = (batch["mask"] > 0) * 1
batch = model(batch)
cur_res = batch[predict_config.out_key][0].permute(1, 2, 0)
cur_res = cur_res.detach().cpu().numpy()
if unpad_to_size is not None:
orig_height, orig_width = unpad_to_size
cur_res = cur_res[:orig_height, :orig_width]
# if unpad_to_size is not None:
# orig_height, orig_width = unpad_to_size
# cur_res = cur_res[:orig_height, :orig_width]
cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
return cur_res
def build_lama_model(config_p: str, ckpt_p: str, device="cuda"):
predict_config = OmegaConf.load(config_p)
predict_config.model.path = ckpt_p
# device = torch.device(predict_config.device)
device = torch.device(device)
train_config_path = os.path.join(predict_config.model.path, "config.yaml")
with open(train_config_path, "r") as f:
train_config = OmegaConf.create(yaml.safe_load(f))
train_config.training_model.predict_only = True
train_config.visualizer.kind = "noop"
checkpoint_path = os.path.join(
predict_config.model.path, "models", predict_config.model.checkpoint
)
model = load_checkpoint(
train_config, checkpoint_path, strict=False, map_location=device
)
model.freeze()
model.to(device)
return model
@torch.no_grad()
def inpaint_img_with_builded_lama(
model, img: np.ndarray, mask: np.ndarray, config_p: str, mod=8, device="cuda"
):
assert len(mask.shape) == 2
if np.max(mask) == 1:
mask = mask * 255
img = torch.from_numpy(img).float().div(255.0)
mask = torch.from_numpy(mask).float()
predict_config = OmegaConf.load(config_p)
batch = {}
batch["image"] = img.permute(2, 0, 1).unsqueeze(0)
batch["mask"] = mask[None, None]
unpad_to_size = [batch["image"].shape[2], batch["image"].shape[3]]
batch["image"] = pad_tensor_to_modulo(batch["image"], mod)
batch["mask"] = pad_tensor_to_modulo(batch["mask"], mod)
batch = move_to_device(batch, device)
batch["mask"] = (batch["mask"] > 0) * 1
batch = model(batch)
cur_res = batch[predict_config.out_key][0].permute(1, 2, 0)
cur_res = cur_res.detach().cpu().numpy()
if unpad_to_size is not None:
orig_height, orig_width = unpad_to_size
cur_res = cur_res[:orig_height, :orig_width]
cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
return cur_res
def setup_args(parser):
parser.add_argument(
"--input_img",
type=str,
required=True,
help="Path to a single input img",
)
parser.add_argument(
"--input_mask_glob",
type=str,
required=True,
help="Glob to input masks",
)
parser.add_argument(
"--output_dir",
type=str,
required=True,
help="Output path to the directory with results.",
)
parser.add_argument(
"--lama_config",
type=str,
default="./third_party/lama/configs/prediction/default.yaml",
help="The path to the config file of lama model. "
"Default: the config of big-lama",
)
parser.add_argument(
"--lama_ckpt",
type=str,
required=True,
help="The path to the lama checkpoint.",
)
if __name__ == "__main__":
"""Example usage:
python lama_inpaint.py \
--input_img FA_demo/FA1_dog.png \
--input_mask_glob "results/FA1_dog/mask*.png" \
--output_dir results \
--lama_config lama/configs/prediction/default.yaml \
--lama_ckpt big-lama
"""
parser = argparse.ArgumentParser()
setup_args(parser)
args = parser.parse_args(sys.argv[1:])
device = "cuda" if torch.cuda.is_available() else "cpu"
img_stem = Path(args.input_img).stem
mask_ps = sorted(glob.glob(args.input_mask_glob))
out_dir = Path(args.output_dir) / img_stem
out_dir.mkdir(parents=True, exist_ok=True)
img = load_img_to_array(args.input_img)
for mask_p in mask_ps:
mask = load_img_to_array(mask_p)
img_inpainted_p = out_dir / f"inpainted_with_{Path(mask_p).name}"
img_inpainted = inpaint_img_with_lama(
img, mask, args.lama_config, args.lama_ckpt, device=device
)
save_array_to_img(img_inpainted, img_inpainted_p)
|