File size: 4,991 Bytes
434b0b0 | 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 | """
Lingteng Qiu
Baseline I2Normal to show the difference demo.
"""
import sys
sys.path.append("./")
import cv2
import einops
import numpy as np
import torch
from tqdm import tqdm
import matplotlib.pyplot as plt
import os
import torch.nn.functional as F
import matplotlib.pyplot as plt
import json
import pdb
import argparse
import tqlt
import random
from tqlt import utils as tu
from tqlt import op as tqlo
from human_generate_system.engineer.NormalEstimator.data_utils import (
HWC3,
resize_image,
norm_normalize,
center_crop,
flip_x,
)
from PIL import Image
from os.path import join
ABS_PATH = join(os.path.dirname(os.path.abspath(__file__)), "DSINE")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="")
parser.add_argument("--num_samples", default=1, type=int)
parser.add_argument("--image_resolution", default=768, type=int)
parser.add_argument("--strength", default=1.0, type=float)
parser.add_argument("--ng_scale", default=1.0, type=float)
parser.add_argument("--ddim_steps", default=10, type=int)
parser.add_argument("--seed", default=23012, type=int)
parser.add_argument("--eta", default=0.0, type=float)
parser.add_argument("--temperature", default=0.0, type=float)
parser.add_argument("--save_memory", action="store_true")
parser.add_argument(
"--negative_prompt",
default="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
type=str,
)
parser.add_argument("--input", "-i", default=None, type=str)
parser.add_argument(
"--prior", default="DSINE", type=str, choices=["DSINE", "geowizard"]
)
parser.add_argument(
"--flip", action="store_true", help="flip init normal and out normal"
)
parser.add_argument("--wo_center", action="store_true", help="without center crop")
parser.add_argument("--num_gpu", default=1, type=int)
parser.add_argument("--rank", default=0, type=int)
opt = parser.parse_args()
assert opt.input is not None and os.path.exists(opt.input)
if os.path.isfile(opt.input):
input_list = [opt.input]
else:
input_name_list = sorted(os.listdir(opt.input))
input_list = [os.path.join(opt.input, name) for name in input_name_list]
all_data = tu.is_img(input_list)
num_samples = opt.num_samples
image_resolution = opt.image_resolution
strength = opt.strength
neg_scale = opt.ng_scale
guess_mode = False
ddim_steps = opt.ddim_steps
seed = opt.seed
eta = opt.eta
temperature = opt.temperature
save_memory = opt.save_memory
tu.seed_everything(seed, verbose=True)
# all_data = sorted(all_data, key=lambda x: x['image'])
if opt.num_gpu > 1:
bucket = len(all_data) // opt.num_gpu
rank = opt.rank
if rank == opt.num_gpu - 1:
all_data = all_data[bucket * rank :]
else:
all_data = all_data[bucket * rank : bucket * (rank + 1)]
if opt.prior == "DSINE":
normal_predictor = torch.hub.load(
ABS_PATH,
"DSINE",
local_file_path="./pretrained_models/dsine.pt",
source="local",
)
else:
raise NotImplementedError
if torch.cuda.is_available():
current_device_id = torch.cuda.current_device()
device = f"cuda:{current_device_id}"
else:
device = "cpu"
output_dir = os.path.join(opt.input, "normal")
os.makedirs(output_dir, exist_ok=True)
for item in tqdm(all_data):
input_image_path = item
basename = os.path.basename(item)
if opt.wo_center:
input_image = cv2.imread(input_image_path)
else:
input_image = center_crop(cv2.imread(input_image_path))
height, width = input_image.shape[:2]
with torch.no_grad():
raw_input_image = HWC3(input_image)
ori_H, ori_W, _ = raw_input_image.shape
img = resize_image(raw_input_image, image_resolution)
H, W, C = img.shape
if opt.prior == "DSINE":
pred_normal = normal_predictor.infer_cv2(img)[0] # (3, H, W)
pred_normal = (pred_normal + 1) / 2 * 255
pred_normal = pred_normal.cpu().numpy().transpose(1, 2, 0)
pred_normal = cv2.cvtColor(
pred_normal.astype(np.uint8), cv2.COLOR_RGB2BGR
)
elif opt.prior == "geowizard":
pred_normal = normal_predictor(img, image_resolution)
pred_normal = (pred_normal + 1) / 2 * 255
pred_normal = cv2.cvtColor(
pred_normal.astype(np.uint8), cv2.COLOR_RGB2BGR
)
pred_normal = cv2.resize(pred_normal, (ori_W, ori_H))
basename = os.path.splitext(basename)[0]
cv2.imwrite(f"{output_dir}/normal_{basename}.png", pred_normal)
|