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
|
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
import os
from collections import OrderedDict
from options.test_options import TestOptions
from options.dataset_options import DatasetOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
import util.util as util
from util.visualizer import Visualizer
from util import html
import torch
from run_engine import run_trt_engine, run_onnx
from datetime import datetime
from PIL import Image, ImageOps
from shutil import copyfile, rmtree
opt = TestOptions().parse(save=False)
data_opt = DatasetOptions().parse(opt.unknown)
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
if data_opt.tag == '':
d = datetime.now()
tag = data_opt.tag = "{}_{}".format(
opt.name,
# opt.experiment,
d.strftime('%Y%m%d%H%M')
)
else:
tag = data_opt.tag
opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/"
print('tag:', tag)
print('render_dir:', render_dir)
util.mkdir(render_dir)
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
start_img_path = os.path.join(render_dir, "frame_00000.png")
copyfile(opt.start_img, start_img_path)
if not opt.engine and not opt.onnx:
model = create_model(opt)
if opt.data_type == 16:
model.half()
elif opt.data_type == 8:
model.type(torch.uint8)
if opt.verbose:
print(model)
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
if opt.data_type == 16:
data['label'] = data['label'].half()
data['inst'] = data['inst'].half()
elif opt.data_type == 8:
data['label'] = data['label'].uint8()
data['inst'] = data['inst'].uint8()
minibatch = 1
print(data['label'])
print(data['inst'])
generated = model.inference(data['label'], data['inst'])
last_path = opt.render_dir + "frame_{:05d}.png".format(i)
tmp_path = opt.render_dir + "frame_{:05d}_tmp.png".format(i+1)
next_path = opt.render_dir + "frame_{:05d}.png".format(i+1)
current_path = opt.render_dir + "ren_{:05d}.png".format(i+1)
print('process image... %s' % last_path)
im = util.tensor2im(generated.data[0])
image_pil = Image.fromarray(im, mode='RGB')
image_pil.save(tmp_path)
os.rename(tmp_path, next_path)
|
