def __init__(self):
self.opt = DemoOptions().parse()
self.opt.data_dir = './datasets/frontend_upload'
self.opt.checkpoints_dir = './checkpoints'
self.model = Pix2PixModel(self.opt)
self.model.opt.inpaint_mode = 'ref'
self.model.eval()
def do_test(opt):
dataloader = data.create_dataloader(opt)
model = Pix2PixModel(opt)
model.eval()
visualizer = Visualizer(opt)
# create a webpage that summarizes the all results
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
# test
for i, data_i in enumerate(dataloader):
if i * opt.batchSize >= opt.how_many:
break
generated = model(data_i, mode='inference')
img_path = data_i['path']
for b in range(generated.shape[0]):
print('process image... %s' % img_path[b])
visuals = OrderedDict([('input_label', data_i['label'][b]),
('synthesized_image', generated[b])])
visualizer.save_images(webpage, visuals, img_path[b:b + 1])
webpage.save()
def __init__(self, opt):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
if len(opt.gpu_ids) > 0:
self.pix2pix_model = DataParallelWithCallback(
self.pix2pix_model, device_ids=opt.gpu_ids)
self.pix2pix_model_on_one_gpu = self.pix2pix_model.module
else:
self.pix2pix_model_on_one_gpu = self.pix2pix_model
self.generated = None
if opt.isTrain:
if not opt.unpairTrain:
(
self.optimizer_G,
self.optimizer_D,
) = self.pix2pix_model_on_one_gpu.create_optimizers(opt)
else:
(
self.optimizer_G,
self.optimizer_D,
self.optimizer_D2,
) = self.pix2pix_model_on_one_gpu.create_optimizers(opt)
self.old_lr = opt.lr
self.d_losses = {}
self.nanCount = 0
def __init__(self, opt, resume_epoch=0):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
if len(opt.gpu_ids) > 1:
self.pix2pix_model = DataParallelWithCallback(
self.pix2pix_model, device_ids=opt.gpu_ids)
self.pix2pix_model_on_one_gpu = self.pix2pix_model.module
else:
self.pix2pix_model.to(opt.gpu_ids[0])
self.pix2pix_model_on_one_gpu = self.pix2pix_model
if opt.use_ema:
self.netG_ema = EMA(opt.ema_beta)
for name, param in self.pix2pix_model_on_one_gpu.net[
'netG'].named_parameters():
if param.requires_grad:
self.netG_ema.register(name, param.data)
self.netCorr_ema = EMA(opt.ema_beta)
for name, param in self.pix2pix_model_on_one_gpu.net[
'netCorr'].named_parameters():
if param.requires_grad:
self.netCorr_ema.register(name, param.data)
self.generated = None
if opt.isTrain:
self.optimizer_G, self.optimizer_D = \
self.pix2pix_model_on_one_gpu.create_optimizers(opt)
self.old_lr = opt.lr
if opt.continue_train and opt.which_epoch == 'latest':
checkpoint = torch.load(
os.path.join(opt.checkpoints_dir, opt.name,
'optimizer.pth'))
self.optimizer_G.load_state_dict(checkpoint['G'])
self.optimizer_D.load_state_dict(checkpoint['D'])
self.last_data, self.last_netCorr, self.last_netG, self.last_optimizer_G = None, None, None, None
def __init__(self, image_q: Queue, done_dict, path):
super(ProcessWorker, self).__init__()
self.image_q = image_q
self.done_dict = done_dict
self.path = path
self.model = Pix2PixModel(opt)
self.model.eval()
self.texture_size = 512
def test_face(old_face_folder,
old_face_label_folder,
name,
gpu_ids,
load_size,
label_nc,
preprocess_mode,
batchSize,
results_dir,
checkpoints_dir="./checkpoints",
how_many=50,
tf_log=False,
tensorboard_log=None,
no_instance=False,
no_parsing_map=False,
serial_batches=False,
nThreads=2,
dataroot="./datasets/cityscapes/",
isTrain=False,
crop_size=512,
aspect_ratio=1.0):
dataloader = data.create_dataloader(batchSize, serial_batches, nThreads,
isTrain, dataroot, old_face_folder,
old_face_label_folder, load_size,
preprocess_mode, crop_size, no_flip,
aspect_ratio)
model = Pix2PixModel(opt)
model.eval()
visualizer = Visualizer(name=name,
checkpoints_dir=checkpoints_dir,
results_dir=results_dir,
batchSize=batchSize,
label_nc=label_nc,
tf_log=tf_log,
tensorboard_log=tensorboard_log)
single_save_url = os.path.join(checkpoints_dir, name, results_dir,
"each_img")
if not os.path.exists(single_save_url):
os.makedirs(single_save_url)
for i, data_i in enumerate(dataloader):
if i * batchSize >= how_many:
break
generated = model(data_i, mode="inference")
img_path = data_i["path"]
for b in range(generated.shape[0]):
img_name = os.path.split(img_path[b])[-1]
save_img_url = os.path.join(single_save_url, img_name)
vutils.save_image((generated[b] + 1) / 2, save_img_url)
def __init__(self, opt):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
self.pix2pix_model_on_one_gpu = self.pix2pix_model
self.generated = None
if opt.isTrain:
self.optimizer_G, self.optimizer_D = \
self.pix2pix_model_on_one_gpu.create_optimizers(opt)
self.old_lr = opt.lr
def __init__(self, opt):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
self.pix2pix_model = torch.nn.parallel.DistributedDataParallel(self.pix2pix_model,
device_ids=[opt.gpu], find_unused_parameters=True)
self.pix2pix_model_on_one_gpu = self.pix2pix_model.module
self.generated = None
if opt.isTrain:
self.optimizer_G, self.optimizer_D = self.pix2pix_model_on_one_gpu.create_optimizers(opt)
self.old_lr = opt.lr
def init(options):
models = {}
models['model'] = None
models['Corr'] = None
models['Gen'] = None
if options.inference_mode == 'pytorch':
models['model'] = Pix2PixModel(options)
models['model'].eval()
else:
models['Corr'] = correspondence_model(options)
models['Gen'] = generate_model(options)
return models
def __init__(self):
opt = TestOptions()
with open(model_config_path, 'r') as fp:
saved_opts = yaml.safe_load(fp)
for _opt in saved_opts:
setattr(opt, _opt, saved_opts[_opt])
model = Pix2PixModel(opt)
model.eval()
self.model = model
# No longer need SPADE-master on PYTHONPATH.
sys.path.pop()
def __init__(self, opt):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
if len(opt.gpu_ids) > 0:
self.pix2pix_model = DataParallelWithCallback(
self.pix2pix_model, device_ids=opt.gpu_ids)
self.pix2pix_model_on_one_gpu = self.pix2pix_model.module
else:
self.pix2pix_model_on_one_gpu = self.pix2pix_model
self.generated = None
if opt.isTrain:
self.optimizer_G, self.optimizer_D = \
self.pix2pix_model_on_one_gpu.create_optimizers(opt)
self.old_lr = opt.lr
def run(verbose=False):
opt = TestOptions().parse(verbose=verbose)
dataset_name = "coco"
dataset_filename = "data." + dataset_name + "_dataset"
datasetlib = importlib.import_module(dataset_filename)
dataset = None
target_dataset_name = dataset_name.replace('_', '') + 'dataset'
for name, cls in datasetlib.__dict__.items():
if name.lower() == target_dataset_name.lower() \
and issubclass(cls, BaseDataset):
dataset = cls
instance = dataset()
instance.initialize(opt)
if verbose:
print("dataset [%s] of size %d was created" %
(type(instance).__name__, len(instance)))
dataloader = torch.utils.data.DataLoader(instance,
batch_size=opt.batchSize,
shuffle=not opt.serial_batches,
num_workers=int(opt.nThreads),
drop_last=opt.isTrain)
model = Pix2PixModel(opt, verbose)
model.eval()
visualizer = Visualizer(opt)
if verbose:
print(dataloader)
for i, data_i in enumerate(dataloader):
if i * opt.batchSize >= opt.how_many:
break
# this is just a dictionary that contains tensors and stuff?
generated = model(data_i, mode='inference', verbose=verbose)
img_path = data_i['path']
for b in range(generated.shape[0]):
# Should only be one?
image_dir = os.path.join(os.path.dirname(__file__), "img")
return visualizer.save_images(img_path[b:b + 1],
generated[b],
image_dir,
verbose=verbose)
def load_model(model_name):
assert model_name in pretrained_models, \
'{} not recongized. Available models: {}'.format(model_name, ', '.join(pretrained_models))
global model, opt, classes
model_subfolder = os.path.join('SPADE/checkpoints/', model_name)
checkpoint_dir = os.path.join(downloads.get_ml4a_downloads_folder(),
model_subfolder)
all_checkpoints_dir = os.path.join(downloads.get_ml4a_downloads_folder(),
'SPADE/checkpoints/')
for gdrive_id, filename in pretrained_models[model_name]:
location = os.path.join(model_subfolder, filename)
downloads.download_from_gdrive(gdrive_id, location)
with open(os.path.join(checkpoint_dir, 'classes_list.txt'),
'r') as classes_file:
classes = eval(classes_file.read())
opt_file = os.path.join(checkpoint_dir, 'opt.txt')
parsed_opt = parse_opt_file(opt_file)
opt = EasyDict({})
opt.isTrain = False
opt.checkpoints_dir = all_checkpoints_dir
opt.name = model_name
opt.aspect_ratio = float(parsed_opt['aspect_ratio'])
opt.load_size = int(parsed_opt['load_size'])
opt.crop_size = int(parsed_opt['crop_size'])
opt.no_instance = True if parsed_opt['no_instance'] == 'True' else False
opt.preprocess_mode = parsed_opt['preprocess_mode']
opt.contain_dontcare_label = True if parsed_opt[
'contain_dontcare_label'] == 'True' else False
opt.gpu_ids = parsed_opt['gpu_ids']
opt.netG = parsed_opt['netG']
opt.ngf = int(parsed_opt['ngf'])
opt.num_upsampling_layers = parsed_opt['num_upsampling_layers']
opt.use_vae = True if parsed_opt['use_vae'] == 'True' else False
opt.label_nc = int(parsed_opt['label_nc'])
opt.semantic_nc = opt.label_nc + (1 if opt.contain_dontcare_label else
0) + (0 if opt.no_instance else 1)
opt.norm_G = parsed_opt['norm_G']
opt.init_type = parsed_opt['init_type']
opt.init_variance = float(parsed_opt['init_variance'])
opt.which_epoch = parsed_opt['which_epoch']
model = Pix2PixModel(opt)
model.eval()
def setup(opts):
global opt
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
options = Options()
parser = options.initialize(parser)
options.parser = parser
name = opts['checkpoints_root'].split('/')[-1]
checkpoints_dir = os.path.join(opts['checkpoints_root'], '..')
parser.set_defaults(name=name)
parser.set_defaults(checkpoints_dir=checkpoints_dir)
opt = options.parse()
opt.name = name
opt.checkpoints_dir = checkpoints_dir
model = Pix2PixModel(opt)
model.eval()
return model
def setup(opts):
global opt
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
options = Options()
parser = options.initialize(parser)
name = 'Labels2Landscapes_512'
checkpoints_dir = os.path.join('.', 'checkpoints')
parser.set_defaults(name=name, checkpoints_dir=checkpoints_dir)
options.parser = parser
opt = options.parse()
opt.name = name
opt.checkpoints_dir = checkpoints_dir
opt.gpu_ids = "0"
device = torch.device('cuda:0')
model = Pix2PixModel(opt)
model.eval()
return model
def __init__(self):
self.opt = TestOptions().parse()
self.opt.crop_size = 512
self.opt.load_size = 512
self.opt.no_instance = True
self.opt.preprocess_mode = "scale_width"
self.opt.dataset_mode = "custom"
self.opt.name = "landscape_pretrained"
self.opt.cache_filelist_read = False
self.opt.cache_filelist_write = False
self.opt.semantic_nc = 182
self.opt.contain_dontcare_label = False
#Load the model
self.model = Pix2PixModel(self.opt)
self.model.eval()
self.map = {
110: 156,
60: 154,
128: 134,
99: 149,
108: 126, #
242: 105,
214: 14,
145: 124,
237: 158,
101: 147, #
70: 96,
135: 168,
65: 148,
209: 143,
150: 110,
83: 125,
120: 135,
141: 119,
156: 161,
82: 177,
230: 118,
113: 123,
232: 162
}
return
def setup(checkpoint_name):
global model, opt, get_params, get_transform, util
spade = eden.setup.get_external_repo_dir('spade')
sys.path.insert(0, spade)
from options.test_options import TestOptions
from models.pix2pix_model import Pix2PixModel
from options.base_options import BaseOptions
from data.base_dataset import get_params, get_transform
import util.util as util
checkpoints_dir = os.path.join(spade, 'checkpoints')
opt_file = os.path.join(os.path.join(checkpoints_dir, checkpoint_name),
'opt.txt')
parsed_opt = parse_opt_file(opt_file)
opt = eden.utils.DictMap()
opt.isTrain = False
opt.checkpoints_dir = checkpoints_dir
opt.name = checkpoint_name
opt.aspect_ratio = float(parsed_opt['aspect_ratio'])
opt.load_size = int(parsed_opt['load_size'])
opt.crop_size = int(parsed_opt['crop_size'])
opt.label_nc = int(parsed_opt['label_nc'])
opt.no_instance = True if parsed_opt['no_instance'] == 'True' else False
opt.preprocess_mode = parsed_opt['preprocess_mode']
opt.contain_dontcare_label = True if parsed_opt[
'contain_dontcare_label'] == 'True' else False
opt.gpu_ids = parsed_opt['gpu_ids']
opt.netG = parsed_opt['netG']
opt.ngf = int(parsed_opt['ngf'])
opt.num_upsampling_layers = parsed_opt['num_upsampling_layers']
opt.use_vae = True if parsed_opt['use_vae'] == 'True' else False
opt.semantic_nc = opt.label_nc + (1 if opt.contain_dontcare_label else
0) + (0 if opt.no_instance else 1)
opt.norm_G = parsed_opt['norm_G']
opt.init_type = parsed_opt['init_type']
opt.init_variance = float(parsed_opt['init_variance'])
opt.which_epoch = parsed_opt['which_epoch']
model = Pix2PixModel(opt)
model.eval()
def main():
opt = TestOptions().parse()
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
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
# read pix2pix/PAN model
if opt.model == 'pix2pix':
assert (opt.dataset_mode == 'aligned')
from models.pix2pix_model import Pix2PixModel
model = Pix2PixModel()
model.initialize(opt)
elif opt.model == 'pan':
from models.pan_model import PanModel
model = PanModel()
model.initialize(opt)
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
# test
for i, data in enumerate(dataset):
if i >= opt.how_many: # default 50 images
break
model.set_input(data)
model.test()
visuals = model.get_current_visuals()
img_path = model.get_image_paths()
print('process image... %s' % img_path)
visualizer.save_images(webpage, visuals, img_path)
webpage.save()
def __init__(self, opt):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
if len(opt.gpu_ids) > 0:
self.pix2pix_model = DataParallelWithCallback(self.pix2pix_model,
device_ids=opt.gpu_ids)
self.pix2pix_model_on_one_gpu = self.pix2pix_model.module
else:
self.pix2pix_model_on_one_gpu = self.pix2pix_model
self.generated = None
if opt.isTrain:
self.optimizer_G, self.optimizer_D = \
self.pix2pix_model_on_one_gpu.create_optimizers(opt)
self.old_lr = opt.lr
self.amp = True if AMP and opt.use_amp else False
if self.amp:
self.scaler_G = GradScaler()
self.scaler_D = GradScaler()
def generate_one(label_img_path, image_path):
opt = TestOptions().parse()
label = Image.open(label_img_path)
# label_np = np.array(label)
params = get_params(opt, label.size)
transform_label = get_transform(opt, params, method=Image.NEAREST, normalize=False)
label_tensor = transform_label(label) * 255.0
label_tensor[label_tensor == 255] = opt.label_nc
label_tensor_batch = torch.unsqueeze(label_tensor,0)
if opt.no_instance:
instance_tensor = 0
image = Image.open(image_path)
image = image.convert('RGB')
transform_image = get_transform(opt, params)
image_tensor = transform_image(image)
input_dict = {'label': label_tensor_batch,
'instance': instance_tensor,
'image': image_tensor,
}
# 加载模型
model = Pix2PixModel(opt)
model.eval()
## 得到3*256*256的图片数据
generated_image = model(input_dict, mode='inference')
generated_image = torch.squeeze(generated_image)
image_numpy = generated_image.numpy()
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
image_numpy = np.clip(image_numpy, 0, 255)
img = image_numpy.astype(np.uint8)
# # 给定image_path即可保存
# image_pil.save(image_path.replace('.jpg', '.png'))
return img
def __init__(self, opt):
self.opt = opt
self.pix2pix_model = Pix2PixModel(opt)
self.pix2pix_model.netG = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
self.pix2pix_model.netG)
self.pix2pix_model.load_networks()
if not opt.isTrain:
self.pix2pix_model.eval()
self.pix2pix_model = self.pix2pix_model.cuda()
if opt.isTrain:
self.optimizer_G, self.optimizer_D = self.pix2pix_model.create_optimizers(
opt)
self.old_lr = opt.lr
self.scaler = amp.GradScaler()
self.local_rank = torch.distributed.get_rank()
self.pix2pix_model = torch.nn.parallel.DistributedDataParallel(
self.pix2pix_model,
device_ids=[self.local_rank],
output_device=self.local_rank,
find_unused_parameters=True)
self.pix2pix_model_on_one_gpu = self.pix2pix_model.module
def run(self):
# Setup
CONFIG = self.deeplab_opt['CONFIG']
self.CONFIG = CONFIG
model_path = self.deeplab_opt['model_path']
cuda = self.deeplab_opt['cuda']
crf = self.deeplab_opt['crf']
camera_id = self.deeplab_opt['camera_id']
device = get_device(cuda)
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
classes = get_classtable(CONFIG)
postprocessor = setup_postprocessor(CONFIG) if crf else None
model = eval(CONFIG.MODEL.NAME)(n_classes=CONFIG.DATASET.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
print("Model:", CONFIG.MODEL.NAME)
# SPADE model
spade_model = Pix2PixModel(self.spade_opt)
spade_model.eval()
print("Spade!")
print(spade_model)
coco_dataset = CocoDataset()
coco_dataset.initialize(self.spade_opt)
print(coco_dataset)
#cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
#np.set_printoptions(threshold=sys.maxsize)
while True:
if len(self.queue) > 0:
frame = self.queue.pop()
print("Original Image shape {}".format(frame.shape))
image, raw_image = preprocessing(frame, device, CONFIG)
print("Image shape {}".format(raw_image.shape))
labelmap = inference(model, image, raw_image, postprocessor)
labelmap[labelmap == 164] = 66
labelmap[labelmap == 109] = 66
#not_dining_mask = ((labelmap < 43) | (labelmap > 50)) & (labelmap != 66) & (labelmap != 0)
#labelmap[not_dining_mask] = 156 # Sky
# tables to dining table
labelmap[labelmap == 66] = 154
dining_objects_mask = (labelmap >= 43) & (labelmap <= 50) | (
labelmap == 138) | (labelmap == 142)
appliances_mask = (labelmap >= 71) & (labelmap <= 79)
labelmap[dining_objects_mask] = 149
labelmap[appliances_mask] = 149
not_sea_mask = (labelmap != 149) & (labelmap != 154)
labelmap[not_sea_mask] = 156
#colormap = self.colorize(labelmap)
#labelmap[labelmap == 66] = 154
# Frisby and more to sea?
#labelmap[labelmap == 33] = 154
#labelmap[labelmap == 66] = 154
#labelmap[labelmap == 80] = 154
#Bottle to flower?
#labelmap[labelmap == 43] = 118
# Person to rock?
#labelmap[labelmap == 0] = 168
#dog to person
#labelmap[labelmap == 17] = 0
# Sky grass and bottle flower
#bottle_mask = (labelmap == 43)
#labelmap[0:193,:] = 156
#labelmap[:,:] = 123
#labelmap[bottle_mask] = 118
#print(labelmap.shape)
#Bottle to potted plant
#labelmap[labelmap == 43] = 63#
#
#dining_stuff = [43,44,45,46,47,48,49,50,66]
#dining_mask = np.isin(labelmap, dining_stuff, invert=True)
#not_dining_mask = (labelmap < 43) | (labelmap > 50) & (labelmap != 66)
#labelmap[dining_objects_mask] = 63
labelimg = Image.fromarray(np.uint8(labelmap), 'L')
label_resized = np.array(
labelimg.resize((256, 256), Image.NEAREST))
uniques = np.unique(labelmap)
instance_counter = 0
instancemap = np.zeros(labelmap.shape)
print(uniques)
for label_id in uniques:
mask = (labelmap == label_id)
instancemap[mask] = instance_counter
instance_counter += 1
instanceimg = Image.fromarray(np.uint8(instancemap), 'L')
#colormap[not_dining_mask] = [0, 0, 0];
item = coco_dataset.get_item_from_images(labelimg, instanceimg)
generated = spade_model(item, mode='inference')
generated_np = util.tensor2im(generated[0])
#color_resized = cv2.cvtColor(np.array(Image.fromarray(colormap).resize((256,256), Image.NEAREST)),cv2.COLOR_BGR2RGB)
generated_np[label_resized == 156] = [0, 0, 0]
# Masking
#print("Generated image shape {} label resize shape {}".format(generated_np.shape, label_resized.shape))
#generated_np[label_resized == 156, :] = [0, 0, 0];
#generated_np[label_resized == 156, :] = [0, 0, 0];
#not_dining_mask = (label_resized < 43) | (label_resized > 50) & (label_resized != 66)
#generated_np[not_dining_mask] = [0, 0, 0];
#only people
#people_mask = np.isin(label_resized, [0,66], invert=True)
#color_resized[label_resized != 0] = [0,0,0]
#generated_rgb = cv2.cvtColor(generated_np, cv2.COLOR_BGR2RGB)
#color_gray = cv2.cvtColor(color_resized, cv2.COLOR_BGR2GRAY)
#color_gray_rgb = cv2.cvtColor(color_gray, cv2.COLOR_GRAY2RGB)
#not_dining_resized = (label_resized < 43) | (label_resized > 50) & (label_resized != 66)
#color_gray_rgb[label_resized != 154, :] = [0, 0, 0];
#generated_np[label_resized == 154, :] = [0,0,0]
#raw_image_resized = cv2.cvtColor(np.array(Image.fromarray(raw_image).resize((256,256), Image.NEAREST)),cv2.COLOR_BGR2RGB)
#raw_image_resized[people_mask] = [0, 0, 0];
#cv2.addWeighted(color_resized, 0.5, raw_image_resized, 0.5 , 0.0, raw_image_resized)
#self.push_frame(raw_image_resized)
#raw_rgb = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)
#final = np.concatenate((generated_np, color_resized), axis=1)
#print("Gans shape {}, colormap shape {}, Final shape {}".format(generated_np.shape, color_resized.shape, final.shape))
#final[:,:256,:] = generated_np
#final[:,256:,:] = color_resized
#raw_rgb = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)
#map_rgb = cv2.cvtColor(colormap, cv2.COLOR_BGR2RGB)
#cv2.addWeighted(map_rgb, 0.5, raw_rgb, 0.5, 0.0, raw_rgb)
#self.push_frame(raw_rgb)
#print("raw image shape {}".format(raw_image.shape))
#print("Generated image {}".format(generated_np))
#print("raw image {}".format(raw_image))
# Register mouse callback function
#cv2.setMouseCallback(window_name, self.mouse_event, labelmap)
# Overlay prediction
#cv2.addWeighted(colormap, 1.0, raw_image, 0.0, 0.0, raw_image)
# Quit by pressing "q" key
#cv2.imshow(window_name, raw_image)
#cv2.resizeWindow(window_name, 1024,1024)
#if cv2.waitKey(10) == ord("q"):
# break
self.push_frame(generated_np)
def live(config_path, model_path, cuda, crf, camera_id):
"""
Inference from camera stream
"""
# Setup
CONFIG = Dict(yaml.load(config_path))
device = get_device(cuda)
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
classes = get_classtable(CONFIG)
postprocessor = setup_postprocessor(CONFIG) if crf else None
model = eval(CONFIG.MODEL.NAME)(n_classes=CONFIG.DATASET.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
print("Model:", CONFIG.MODEL.NAME)
# SPADE model
opt = TestOptions().parse()
opt.use_vae = False
spade_model = Pix2PixModel(opt)
spade_model.eval()
print("Spade!")
print(spade_model)
coco_dataset = CocoDataset()
coco_dataset.initialize(opt)
print(coco_dataset)
# UVC camera stream
cap = cv2.VideoCapture(camera_id)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"YUYV"))
def colorize(labelmap):
print(labelmap.shape)
# Assign a unique color to each label
labelmap = labelmap.astype(np.float32) / CONFIG.DATASET.N_CLASSES
colormap = cm.jet_r(labelmap)[..., :-1] * 255.0
return np.uint8(colormap)
def mouse_event(event, x, y, flags, labelmap):
# Show a class name of a mouse-overed pixel
label = labelmap[y, x]
name = classes[label]
print(name)
window_name = "{} + {}".format(CONFIG.MODEL.NAME, CONFIG.DATASET.NAME)
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
np.set_printoptions(threshold=sys.maxsize)
while True:
_, frame = cap.read()
image, raw_image = preprocessing(frame, device, CONFIG)
#print("Image shape {}".format(image.shape))
labelmap = inference(model, image, raw_image, postprocessor)
# Frisby and more to sea?
#labelmap[labelmap == 33] = 154
#labelmap[labelmap == 66] = 154
#labelmap[labelmap == 80] = 154
#Bottle to flower?
labelmap[labelmap == 43] = 118
# Person to rock?
#labelmap[labelmap == 0] = 168
#dog to person
#labelmap[labelmap == 17] = 0
# Sky grass and bottle flower
#bottle_mask = (labelmap == 43)
#labelmap[0:193,:] = 156
#labelmap[:,:] = 123
#labelmap[bottle_mask] = 118
#print(labelmap.shape)
#colormap = colorize(labelmap)
uniques = np.unique(labelmap)
instance_counter = 0
instancemap = np.zeros(labelmap.shape)
print(uniques)
for label_id in uniques:
mask = (labelmap == label_id)
instancemap[mask] = instance_counter
instance_counter += 1
labelimg = Image.fromarray(np.uint8(labelmap), 'L')
instanceimg = Image.fromarray(np.uint8(instancemap), 'L')
#labelimg.show()
item = coco_dataset.get_item_from_images(labelimg, instanceimg)
generated = spade_model(item, mode='inference')
generated_np = util.tensor2im(generated[0])
# Masking
#print("Generated image shape {} label resize shape {}".format(generated_np.shape, label_resized.shape))
#label_resized = np.array(labelimg.resize((256,256), Image.NEAREST))
#generated_np[label_resized != 118, :] = [0, 0, 0];
generated_rgb = cv2.cvtColor(generated_np, cv2.COLOR_BGR2RGB)
#print("raw image shape {}".format(raw_image.shape))
#print("Generated image {}".format(generated_np))
#print("raw image {}".format(raw_image))
# Register mouse callback function
cv2.setMouseCallback(window_name, mouse_event, labelmap)
# Overlay prediction
#cv2.addWeighted(colormap, 1.0, raw_image, 0.0, 0.0, raw_image)
# Quit by pressing "q" key
cv2.imshow(window_name, generated_rgb)
cv2.resizeWindow(window_name, 1024, 1024)
if cv2.waitKey(10) == ord("q"):
break
elif self.menu.doFill:
sub_pos = (event.pos[0] - self.MENU_WIDTH, event.pos[1]
) # 转换为subsurface中pos
self.fill_by_cv2(sub_pos)
else:
sub_pos = (event.pos[0] - self.MENU_WIDTH, event.pos[1]
) # 转换为subsurface中pos
self.brush.start_draw(sub_pos) # subsurface适配
elif event.type == MOUSEMOTION:
sub_pos = (event.pos[0] - self.MENU_WIDTH, event.pos[1]
) # 转换为subsurface中pos
self.brush.draw(sub_pos) # subsurface适配
elif event.type == MOUSEBUTTONUP:
self.brush.end_draw()
self.menu.draw()
pygame.display.update()
def main():
app = Painter()
app.run()
if __name__ == '__main__':
# 加载模型
model = Pix2PixModel(TestOptions().parse())
model.eval()
main()
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import TensorBoardLogger
from models.pix2pix_model import Pix2PixModel
if __name__ == "__main__":
logger = TensorBoardLogger('tb_logs', name='my_model')
trainer = Trainer(logger=logger)
net = Pix2PixModel(input_nc=3,
output_nc=3,
ngf=64,
ndf=64,
netG='unet_128',
gan_mode='vanilla',
norm='instance',
n_layers_D=3,
netD='n_layers',
lr=0.0002,
beta1=0.5,
device='cpu',
batch_size=3,
dataroot='/home/gujiaxin/桌面/AB',
load_size=(300, 300),
crop_size=(300, 300),
preprocess=['resize', 'crop', 'flip'],
lambda_L1=100.0)
trainer.fit(net)
from util import html
opt = TestOptions().parse()
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
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
# read pix2pix/PAN model
if opt.model == 'pix2pix':
assert (opt.dataset_mode == 'aligned')
from models.pix2pix_model import Pix2PixModel
model = Pix2PixModel()
model.initialize(opt)
elif opt.model == 'pan':
from models.pan_model import PanModel
model = PanModel()
model.initialize(opt)
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
def run(self):
# Setup
CONFIG = self.deeplab_opt['CONFIG']
self.CONFIG = CONFIG
model_path = self.deeplab_opt['model_path']
cuda = self.deeplab_opt['cuda']
crf = self.deeplab_opt['crf']
camera_id = self.deeplab_opt['camera_id']
device = get_device(cuda)
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
classes = get_classtable(CONFIG)
postprocessor = setup_postprocessor(CONFIG) if crf else None
model = eval(CONFIG.MODEL.NAME)(n_classes=CONFIG.DATASET.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
print("Model:", CONFIG.MODEL.NAME)
# SPADE model
spade_model = Pix2PixModel(self.spade_opt)
spade_model.eval()
spade_model.to(device)
print("Spade!")
print(spade_model)
coco_dataset = CocoDataset()
coco_dataset.initialize(self.spade_opt)
print(coco_dataset)
while True:
while not self.osc_queue.empty():
item = self.osc_queue.get()
self.process_queue(item)
if len(self.queue) > 0:
frame = self.queue.pop()
#print("Original Image shape {}".format(frame.shape))
image, raw_image = preprocessing(frame, device, CONFIG)
raw_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)
#print("Image shape {}".format(raw_image.shape))
labelmap = inference(model, image, raw_image, postprocessor)
if self.current_state == "test-bowl":
self.test_bowl(labelmap)
uniques = np.unique(labelmap)
print([ID_TO_LABEL[unique] for unique in uniques])
if self.send_bowl and LABEL_TO_ID['bowl'] in uniques:
box = self.get_bounding_box_of(LABEL_TO_ID['bowl'],
labelmap)
(rmin, cmin, rmax, cmax) = box
coords = (int((cmin + cmax) / 2), int((rmin + rmax) / 2))
print("Bowl coords {} out of {} ".format(
coords, labelmap.shape))
self.t2i_client.send_message("/deeplab/bowl",
[coords[0], coords[1]])
if not self.map_deeplab:
colormap = self.colorize(labelmap)
for masking in self.deeplab_masks:
mask = np.isin(labelmap,
masking['items'],
invert=masking['invert'])
colormap[mask, :] = [0, 0, 0]
for mapping in self.maps:
mask = np.isin(labelmap,
mapping['from'],
invert=mapping['invert'])
if mapping['expand'] > 0:
mask = self.expand_mask(mask, mapping['expand'])
labelmap[mask] = mapping['to']
if self.map_deeplab:
colormap = self.colorize(labelmap)
for masking in self.deeplab_masks:
mask = np.isin(labelmap,
masking['items'],
invert=masking['invert'])
colormap[mask, :] = [0, 0, 0]
if self.show_raw:
for masking in self.deeplab_masks:
mask = np.isin(labelmap,
masking['items'],
invert=masking['invert'])
raw_image[mask, :] = [0, 0, 0]
if self.show_labels:
for unique in uniques:
box = self.get_bounding_box_of(unique, labelmap)
self.put_text_in_center(colormap, box,
ID_TO_LABEL[unique])
#color_resized = cv2.cvtColor(np.array(Image.fromarray(colormap).resize((256,256), Image.NEAREST)),cv2.COLOR_BGR2RGB)
if self.show_gaugan:
uniques = np.unique(labelmap)
instance_counter = 0
instancemap = np.zeros(labelmap.shape)
for label_id in uniques:
mask = (labelmap == label_id)
instancemap[mask] = instance_counter
instance_counter += 1
instanceimg = Image.fromarray(np.uint8(instancemap), 'L')
labelimg = Image.fromarray(np.uint8(labelmap), 'L')
label_resized = np.array(
labelimg.resize((256, 256), Image.NEAREST))
item = coco_dataset.get_item_from_images(
labelimg, instanceimg)
generated = spade_model(item, mode='inference')
generated_np = util.tensor2im(generated[0])
for masking in self.gaugan_masks:
mask = np.isin(label_resized,
masking['items'],
invert=masking['invert'])
generated_np[mask, :] = [0, 0, 0]
print("SPADE Shape {}".format(generated_np.shape))
else:
generated_np = np.uint8(np.zeros((256, 256, 3)))
final = np.concatenate((generated_np, colormap, raw_image),
axis=1)
#final = np.concatenate((generated_np, colormap), axis=1)
self.push_frame(final)
def __init__(self, opt):
self.model = Pix2PixModel(opt)
self.model.eval()
import os
from collections import OrderedDict
import torch
import data
from options.test_options import TestOptions
from models.pix2pix_model import Pix2PixModel
from util.visualizer import Visualizer
from util import html
opt = TestOptions().parse()
dataloader = data.create_dataloader(opt)
model = Pix2PixModel(opt)
model.eval()
visualizer = Visualizer(opt)
# create a webpage that summarizes the all results
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
# test
for i, data_i in enumerate(dataloader):
if i * opt.batchSize >= opt.how_many:
break
def __init__(self):
self.model = Pix2PixModel(opt)
self.model.eval()
torch.cuda.set_device(3)
self.cuda = torch.cuda.is_available()
