3DViewOpt/main.py at main · tai5863/3DViewOpt · GitHub

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
from core.result_manager import ResultManager
from core.data import get_data, filter_unique_files, generate_dpt_depth
from Optuna.optimize_optuna import optimize_optuna
from PyTorch.optimize_pytorch import optimize_pytorch
import os
import config
import torch.utils.data.distributed
from third_party._3d_moments.core.inpainter import Inpainter
from core.renderer import Renderer
from core.iaa_processor import IAAModelManager
import glob

from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True


def initialize(args):
    """
    Initialize the process.
    """

    device = "cuda:{}".format(args.local_rank)

    inpainter = Inpainter(args)
    renderer = Renderer(args, inpainter, device)
    iaa_model = IAAModelManager(args, device)

    if not os.path.exists(args.data_dir):
        raise ValueError(f'{args.data_dir} does not exist.')

    # get outpainted images
    img_dir = os.path.join(args.data_dir, 'input')
    outpainted_img_dir = os.path.join(args.data_dir, 'outpainted')

    # generate DPT depth maps
    if args.dpt_depth:
        generate_dpt_depth(outpainted_img_dir, os.path.join(
            args.data_dir, 'dpt_depth'))

    # get image files
    img_files = sorted(glob.glob(os.path.join(img_dir, '*.png')) +
                       glob.glob(os.path.join(img_dir, '*.jpg')))

    outpainted_img_files = sorted(glob.glob(os.path.join(outpainted_img_dir, '*.png')) +
                                  glob.glob(os.path.join(outpainted_img_dir, '*.jpg')))

    # filter unique files
    img_files = filter_unique_files(img_files, outpainted_img_files)
    outpainted_img_files = filter_unique_files(outpainted_img_files, img_files)

    # sort files
    img_files = sorted(img_files)
    outpainted_img_files = sorted(outpainted_img_files)

    return device, renderer, iaa_model, img_files, outpainted_img_files


def run(args):
    device, renderer, iaa_model, img_files, outpainted_img_files = initialize(args)

    out_folder = os.path.join(os.path.join(args.data_dir, args.output_dir))
    os.makedirs(out_folder, exist_ok=True)

    lr = args.lr
    iteration = args.iter
    w_mask = args.w_mask
    max_image_num = args.max_image_num
    optimization = args.optimization

    print('{}: lr: {}, iteration: {}, w_mask: {}'.format(optimization, lr, iteration, w_mask))

    print('Executing optimization...')

    for img_id in range(len(img_files)):

        if (img_id + 1) > max_image_num:
            break

        # Get the image files
        img_file = img_files[img_id]
        outpainted_img_file = outpainted_img_files[img_id]
        base_name = os.path.basename(outpainted_img_file).split(".")[0]
        dpt_depth_file = os.path.join(args.data_dir, 'dpt_depth', f'{base_name}.png')

        data = get_data(img_file, outpainted_img_file, dpt_depth_file, device)

        print(f'Processing {os.path.basename(img_file)}')

        # Get the target image points and rgb
        with torch.no_grad():
            pts, rgba, _ = renderer.start_process(data)

        cam_poses = None

        # Execute the optimization
        if optimization == 'cmaes':
            saver = ResultManager(args, data, out_folder, optimization, saving_option={'gif': False, 'video': False, 'tensorboard': False, 'logs': False})
            optimize_optuna(data, iteration, w_mask, pts, rgba, renderer, iaa_model, out_folder, saver, device)

        elif optimization == 'gradient':
            saver = ResultManager(args, data, out_folder, optimization, saving_option={'gif': False, 'video': False, 'tensorboard': False, 'logs': False})
            optimize_pytorch(data, lr, iteration, w_mask, pts, rgba, cam_poses, renderer, iaa_model, saver, device, from_random_pose=True)

        else:
            assert False, 'Invalid optimization method'


if __name__ == '__main__':
    args = config.config_parser()

    run(args)

    print("Done!!!!!")