train.py

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()