import argparse
import os
import time
from datetime import datetime
from distutils.util import strtobool
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from data_loader import (FileDataset,
RandomResizedCropWithAutoCenteringAndZeroPadding)
from torch.utils.data.distributed import DistributedSampler
from conr import CoNR
from tqdm import tqdm
def data_sampler(dataset, shuffle, distributed):
if distributed:
return torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)
if shuffle:
return torch.utils.data.RandomSampler(dataset)
else:
return torch.utils.data.SequentialSampler(dataset)
def save_output(image_name, inputs_v, d_dir=".", crop=None):
import cv2
inputs_v = inputs_v.detach().squeeze()
input_np = torch.clamp(inputs_v*255, 0, 255).byte().cpu().numpy().transpose(
(1, 2, 0))
# cv2.setNumThreads(1)
out_render_scale = cv2.cvtColor(input_np, cv2.COLOR_RGBA2BGRA)
if crop is not None:
crop = crop.cpu().numpy()[0]
output_img = np.zeros((crop[0], crop[1], 4), dtype=np.uint8)
before_resize_scale = cv2.resize(
out_render_scale, (crop[5]-crop[4]+crop[8]+crop[9], crop[3]-crop[2]+crop[6]+crop[7]), interpolation=cv2.INTER_AREA) # w,h
output_img[crop[2]:crop[3], crop[4]:crop[5]] = before_resize_scale[crop[6]:before_resize_scale.shape[0] -
crop[7], crop[8]:before_resize_scale.shape[1]-crop[9]]
else:
output_img = out_render_scale
cv2.imwrite(d_dir+"/"+image_name.split(os.sep)[-1]+'.png',
output_img
)
def test():
source_names_list = []
for name in sorted(os.listdir(args.test_input_person_images)):
thissource = os.path.join(args.test_input_person_images, name)
if os.path.isfile(thissource):
source_names_list.append(thissource)
if os.path.isdir(thissource):
print("skipping empty folder :"+thissource)
image_names_list = []
for name in sorted(os.listdir(args.test_input_poses_images)):
thistarget = os.path.join(args.test_input_poses_images, name)
if os.path.isfile(thistarget):
image_names_list.append([thistarget, *source_names_list])
if os.path.isdir(thistarget):
print("skipping folder :"+thistarget)
print(image_names_list)
print("---building models")
conrmodel = CoNR(args)
conrmodel.load_model(path=args.test_checkpoint_dir)
conrmodel.dist()
infer(args, conrmodel, image_names_list)
def infer(args, humanflowmodel, image_names_list):
print("---test images: ", len(image_names_list))
test_salobj_dataset = FileDataset(image_names_list=image_names_list,
fg_img_lbl_transform=transforms.Compose([
RandomResizedCropWithAutoCenteringAndZeroPadding(
(args.dataloader_imgsize, args.dataloader_imgsize), scale=(1, 1), ratio=(1.0, 1.0), center_jitter=(0.0, 0.0)
)]),
shader_pose_use_gt_udp_test=not args.test_pose_use_parser_udp,
shader_target_use_gt_rgb_debug=False
)
sampler = data_sampler(test_salobj_dataset, shuffle=False,
distributed=args.distributed)
train_data = DataLoader(test_salobj_dataset,
batch_size=1,
shuffle=False,sampler=sampler,
num_workers=args.dataloaders)
# start testing
train_num = train_data.__len__()
time_stamp = time.time()
prev_frame_rgb = []
prev_frame_a = []
pbar = tqdm(range(train_num), ncols=100)
for i, data in enumerate(train_data):
data_time_interval = time.time() - time_stamp
time_stamp = time.time()
with torch.no_grad():
data["character_images"] = torch.cat(
[data["character_images"], *prev_frame_rgb], dim=1)
data["character_masks"] = torch.cat(
[data["character_masks"], *prev_frame_a], dim=1)
data = humanflowmodel.data_norm_image(data)
pred = humanflowmodel.model_step(data, training=False)
# remember to call humanflowmodel.reset_charactersheet() if you change character .
train_time_interval = time.time() - time_stamp
time_stamp = time.time()
if args.local_rank == 0:
pbar.set_description(f"Epoch {i}/{train_num}")
pbar.set_postfix({"data_time": data_time_interval, "train_time":train_time_interval})
pbar.update(1)
with torch.no_grad():
if args.test_output_video:
pred_img = pred["shader"]["y_weighted_warp_decoded_rgba"]
save_output(
str(int(data["imidx"].cpu().item())), pred_img, args.test_output_dir, crop=data["pose_crop"])
if args.test_output_udp:
pred_img = pred["shader"]["x_target_sudp_a"]
save_output(
"udp_"+str(int(data["imidx"].cpu().item())), pred_img, args.test_output_dir)
def build_args():
parser = argparse.ArgumentParser()
# distributed learning settings
parser.add_argument("--world_size", type=int, default=1,
help='world size')
parser.add_argument("--local_rank", type=int, default=0,
help='local_rank, DON\'T change it')
# model settings
parser.add_argument('--dataloader_imgsize', type=int, default=256,
help='Input image size of the model')
parser.add_argument('--batch_size', type=int, default=4,
help='minibatch size')
parser.add_argument('--model_name', default='model_result',
help='Name of the experiment')
parser.add_argument('--dataloaders', type=int, default=2,
help='Num of dataloaders')
parser.add_argument('--mode', default="test", choices=['train', 'test'],
help='Training mode or Testing mode')
# i/o settings
parser.add_argument('--test_input_person_images',
type=str, default="./character_sheet/",
help='Directory to input character sheets')
parser.add_argument('--test_input_poses_images', type=str,
default="./test_data/",
help='Directory to input UDP sequences or pose images')
parser.add_argument('--test_checkpoint_dir', type=str,
default='./weights/',
help='Directory to model weights')
parser.add_argument('--test_output_dir', type=str,
default="./results/",
help='Directory to output images')
# output content settings
parser.add_argument('--test_output_video', type=strtobool, default=True,
help='Whether to output the final result of CoNR, \
images will be output to test_output_dir while True.')
parser.add_argument('--test_output_udp', type=strtobool, default=False,
help='Whether to output UDP generated from UDP detector, \
this is meaningful ONLY when test_input_poses_images \
is not UDP sequences but pose images. Meanwhile, \
test_pose_use_parser_udp need to be True')
# UDP detector settings
parser.add_argument('--test_pose_use_parser_udp',
type=strtobool, default=False,
help='Whether to use UDP detector to generate UDP from pngs, \
pose input MUST be pose images instead of UDP sequences \
while True')
args = parser.parse_args()
args.distributed = (args.world_size > 1)
if args.local_rank == 0:
print("batch_size:", args.batch_size, flush=True)
if args.distributed:
if args.local_rank == 0:
print("world_size: ", args.world_size)
torch.distributed.init_process_group(
backend="nccl", init_method="env://", world_size=args.world_size)
torch.cuda.set_device(args.local_rank)
torch.backends.cudnn.benchmark = True
else:
args.local_rank = 0
return args
if __name__ == "__main__":
args = build_args()
test()