From cde29be4f8c4fcd316d701a147fd8679fc7209d9 Mon Sep 17 00:00:00 2001 From: hypox64 Date: Wed, 14 Aug 2019 15:11:44 +0800 Subject: [PATCH] support for pix2pixHD model --- deepmosaic.py | 9 +- models/loadmodel.py | 8 +- models/pix2pixHD_model.py | 418 ++++++++++++++++++++++++++++++++++++++ models/runmodel.py | 11 +- models/unet_parts.py | 2 +- options.py | 11 +- util/data.py | 3 +- util/mosaic.py | 18 +- util/util.py | 8 + 9 files changed, 467 insertions(+), 21 deletions(-) create mode 100644 models/pix2pixHD_model.py diff --git a/deepmosaic.py b/deepmosaic.py index 63583cb..fc78e25 100644 --- a/deepmosaic.py +++ b/deepmosaic.py @@ -42,7 +42,7 @@ if opt.mode == 'add': positions.append([x,y,area]) cv2.imwrite(os.path.join('./tmp/ROI_mask', os.path.basename(imagepath)),mask) - print('Optimized ROI locations...') + print('Optimize ROI locations...') mask_index = filt.position_medfilt(np.array(positions), 7) # add mosaic @@ -71,7 +71,7 @@ elif opt.mode == 'clean': img_result = img_origin.copy() if size != 0 : img_mosaic = img_origin[y-size:y+size,x-size:x+size] - img_fake=runmodel.run_pix2pix(img_mosaic,netG,use_gpu = opt.use_gpu) + img_fake=runmodel.run_pix2pix(img_mosaic,netG,opt) img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather) cv2.imwrite(os.path.join(opt.result_dir,os.path.basename(path)),img_result) @@ -90,7 +90,8 @@ elif opt.mode == 'clean': img_origin = impro.imread(imagepath) x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt) positions.append([x,y,size]) - print('Find Positions:',imagepath) + print('Find mosaic location:',imagepath) + print('Optimize mosaic locations...') positions =np.array(positions) for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num) @@ -102,7 +103,7 @@ elif opt.mode == 'clean': img_result = img_origin.copy() if size != 0: img_mosaic = img_origin[y-size:y+size,x-size:x+size] - img_fake = runmodel.run_pix2pix(img_mosaic,netG,use_gpu = opt.use_gpu) + img_fake = runmodel.run_pix2pix(img_mosaic,netG,opt) img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather) cv2.imwrite(os.path.join('./tmp/replace_mosaic',os.path.basename(imagepath)),img_result) print('Clean Mosaic:',imagepath) diff --git a/models/loadmodel.py b/models/loadmodel.py index eaeb06b..0dd65ce 100755 --- a/models/loadmodel.py +++ b/models/loadmodel.py @@ -1,10 +1,14 @@ import torch -from .pix2pix_model import * +from .pix2pix_model import define_G +from .pix2pixHD_model import define_G as define_G_HD from .unet_model import UNet def pix2pix(opt): # print(opt.model_path,opt.netG) - netG = define_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[]) + if opt.netG == 'HD': + netG = define_G_HD(3, 3, 64, 'global' ,4) + else: + netG = define_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[]) netG.load_state_dict(torch.load(opt.model_path)) netG.eval() diff --git a/models/pix2pixHD_model.py b/models/pix2pixHD_model.py new file mode 100644 index 0000000..b20369b --- /dev/null +++ b/models/pix2pixHD_model.py @@ -0,0 +1,418 @@ +# This code clone from https://github.com/NVIDIA/pix2pixHD +# LICENSE file : https://github.com/NVIDIA/pix2pixHD/blob/master/LICENSE.txt +import torch +import torch.nn as nn +import functools +from torch.autograd import Variable +import numpy as np + +############################################################################### +# Functions +############################################################################### +def weights_init(m): + classname = m.__class__.__name__ + if classname.find('Conv') != -1: + m.weight.data.normal_(0.0, 0.02) + elif classname.find('BatchNorm2d') != -1: + m.weight.data.normal_(1.0, 0.02) + m.bias.data.fill_(0) + +def get_norm_layer(norm_type='instance'): + if norm_type == 'batch': + norm_layer = functools.partial(nn.BatchNorm2d, affine=True) + elif norm_type == 'instance': + norm_layer = functools.partial(nn.InstanceNorm2d, affine=False) + else: + raise NotImplementedError('normalization layer [%s] is not found' % norm_type) + return norm_layer + +def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1, + n_blocks_local=3, norm='instance', gpu_ids=[]): + norm_layer = get_norm_layer(norm_type=norm) + if netG == 'global': + netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer) + elif netG == 'local': + netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, + n_local_enhancers, n_blocks_local, norm_layer) + elif netG == 'encoder': + netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, norm_layer) + else: + raise('generator not implemented!') + # print(netG) + if len(gpu_ids) > 0: + assert(torch.cuda.is_available()) + netG.cuda(gpu_ids[0]) + netG.apply(weights_init) + return netG + +def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False, gpu_ids=[]): + norm_layer = get_norm_layer(norm_type=norm) + netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat) + print(netD) + if len(gpu_ids) > 0: + assert(torch.cuda.is_available()) + netD.cuda(gpu_ids[0]) + netD.apply(weights_init) + return netD + +def print_network(net): + if isinstance(net, list): + net = net[0] + num_params = 0 + for param in net.parameters(): + num_params += param.numel() + print(net) + print('Total number of parameters: %d' % num_params) + +############################################################################## +# Losses +############################################################################## +class GANLoss(nn.Module): + def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0, + tensor=torch.FloatTensor): + super(GANLoss, self).__init__() + self.real_label = target_real_label + self.fake_label = target_fake_label + self.real_label_var = None + self.fake_label_var = None + self.Tensor = tensor + if use_lsgan: + self.loss = nn.MSELoss() + else: + self.loss = nn.BCELoss() + + def get_target_tensor(self, input, target_is_real): + target_tensor = None + if target_is_real: + create_label = ((self.real_label_var is None) or + (self.real_label_var.numel() != input.numel())) + if create_label: + real_tensor = self.Tensor(input.size()).fill_(self.real_label) + self.real_label_var = Variable(real_tensor, requires_grad=False) + target_tensor = self.real_label_var + else: + create_label = ((self.fake_label_var is None) or + (self.fake_label_var.numel() != input.numel())) + if create_label: + fake_tensor = self.Tensor(input.size()).fill_(self.fake_label) + self.fake_label_var = Variable(fake_tensor, requires_grad=False) + target_tensor = self.fake_label_var + return target_tensor + + def __call__(self, input, target_is_real): + if isinstance(input[0], list): + loss = 0 + for input_i in input: + pred = input_i[-1] + target_tensor = self.get_target_tensor(pred, target_is_real) + loss += self.loss(pred, target_tensor) + return loss + else: + target_tensor = self.get_target_tensor(input[-1], target_is_real) + return self.loss(input[-1], target_tensor) + +class VGGLoss(nn.Module): + def __init__(self, gpu_ids): + super(VGGLoss, self).__init__() + self.vgg = Vgg19().cuda() + self.criterion = nn.L1Loss() + self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0] + + def forward(self, x, y): + x_vgg, y_vgg = self.vgg(x), self.vgg(y) + loss = 0 + for i in range(len(x_vgg)): + loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach()) + return loss + +############################################################################## +# Generator +############################################################################## +class LocalEnhancer(nn.Module): + def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9, + n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'): + super(LocalEnhancer, self).__init__() + self.n_local_enhancers = n_local_enhancers + + ###### global generator model ##### + ngf_global = ngf * (2**n_local_enhancers) + model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model + model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers + self.model = nn.Sequential(*model_global) + + ###### local enhancer layers ##### + for n in range(1, n_local_enhancers+1): + ### downsample + ngf_global = ngf * (2**(n_local_enhancers-n)) + model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0), + norm_layer(ngf_global), nn.ReLU(True), + nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1), + norm_layer(ngf_global * 2), nn.ReLU(True)] + ### residual blocks + model_upsample = [] + for i in range(n_blocks_local): + model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer)] + + ### upsample + model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1), + norm_layer(ngf_global), nn.ReLU(True)] + + ### final convolution + if n == n_local_enhancers: + model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()] + + setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample)) + setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample)) + + self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False) + + def forward(self, input): + ### create input pyramid + input_downsampled = [input] + for i in range(self.n_local_enhancers): + input_downsampled.append(self.downsample(input_downsampled[-1])) + + ### output at coarest level + output_prev = self.model(input_downsampled[-1]) + ### build up one layer at a time + for n_local_enhancers in range(1, self.n_local_enhancers+1): + model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1') + model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2') + input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers] + output_prev = model_upsample(model_downsample(input_i) + output_prev) + return output_prev + +class GlobalGenerator(nn.Module): + def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d, + padding_type='reflect'): + assert(n_blocks >= 0) + super(GlobalGenerator, self).__init__() + activation = nn.ReLU(True) + + model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation] + ### downsample + for i in range(n_downsampling): + mult = 2**i + model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), + norm_layer(ngf * mult * 2), activation] + + ### resnet blocks + mult = 2**n_downsampling + for i in range(n_blocks): + model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)] + + ### upsample + for i in range(n_downsampling): + mult = 2**(n_downsampling - i) + model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), + norm_layer(int(ngf * mult / 2)), activation] + model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()] + self.model = nn.Sequential(*model) + + def forward(self, input): + return self.model(input) + +# Define a resnet block +class ResnetBlock(nn.Module): + def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False): + super(ResnetBlock, self).__init__() + self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout) + + def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout): + conv_block = [] + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + + conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), + norm_layer(dim), + activation] + if use_dropout: + conv_block += [nn.Dropout(0.5)] + + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), + norm_layer(dim)] + + return nn.Sequential(*conv_block) + + def forward(self, x): + out = x + self.conv_block(x) + return out + +class Encoder(nn.Module): + def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d): + super(Encoder, self).__init__() + self.output_nc = output_nc + + model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), + norm_layer(ngf), nn.ReLU(True)] + ### downsample + for i in range(n_downsampling): + mult = 2**i + model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), + norm_layer(ngf * mult * 2), nn.ReLU(True)] + + ### upsample + for i in range(n_downsampling): + mult = 2**(n_downsampling - i) + model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), + norm_layer(int(ngf * mult / 2)), nn.ReLU(True)] + + model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()] + self.model = nn.Sequential(*model) + + def forward(self, input, inst): + outputs = self.model(input) + + # instance-wise average pooling + outputs_mean = outputs.clone() + inst_list = np.unique(inst.cpu().numpy().astype(int)) + for i in inst_list: + for b in range(input.size()[0]): + indices = (inst[b:b+1] == int(i)).nonzero() # n x 4 + for j in range(self.output_nc): + output_ins = outputs[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] + mean_feat = torch.mean(output_ins).expand_as(output_ins) + outputs_mean[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat + return outputs_mean + +class MultiscaleDiscriminator(nn.Module): + def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, + use_sigmoid=False, num_D=3, getIntermFeat=False): + super(MultiscaleDiscriminator, self).__init__() + self.num_D = num_D + self.n_layers = n_layers + self.getIntermFeat = getIntermFeat + + for i in range(num_D): + netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat) + if getIntermFeat: + for j in range(n_layers+2): + setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j))) + else: + setattr(self, 'layer'+str(i), netD.model) + + self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False) + + def singleD_forward(self, model, input): + if self.getIntermFeat: + result = [input] + for i in range(len(model)): + result.append(model[i](result[-1])) + return result[1:] + else: + return [model(input)] + + def forward(self, input): + num_D = self.num_D + result = [] + input_downsampled = input + for i in range(num_D): + if self.getIntermFeat: + model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)] + else: + model = getattr(self, 'layer'+str(num_D-1-i)) + result.append(self.singleD_forward(model, input_downsampled)) + if i != (num_D-1): + input_downsampled = self.downsample(input_downsampled) + return result + +# Defines the PatchGAN discriminator with the specified arguments. +class NLayerDiscriminator(nn.Module): + def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False): + super(NLayerDiscriminator, self).__init__() + self.getIntermFeat = getIntermFeat + self.n_layers = n_layers + + kw = 4 + padw = int(np.ceil((kw-1.0)/2)) + sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]] + + nf = ndf + for n in range(1, n_layers): + nf_prev = nf + nf = min(nf * 2, 512) + sequence += [[ + nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw), + norm_layer(nf), nn.LeakyReLU(0.2, True) + ]] + + nf_prev = nf + nf = min(nf * 2, 512) + sequence += [[ + nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw), + norm_layer(nf), + nn.LeakyReLU(0.2, True) + ]] + + sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]] + + if use_sigmoid: + sequence += [[nn.Sigmoid()]] + + if getIntermFeat: + for n in range(len(sequence)): + setattr(self, 'model'+str(n), nn.Sequential(*sequence[n])) + else: + sequence_stream = [] + for n in range(len(sequence)): + sequence_stream += sequence[n] + self.model = nn.Sequential(*sequence_stream) + + def forward(self, input): + if self.getIntermFeat: + res = [input] + for n in range(self.n_layers+2): + model = getattr(self, 'model'+str(n)) + res.append(model(res[-1])) + return res[1:] + else: + return self.model(input) + +from torchvision import models +class Vgg19(torch.nn.Module): + def __init__(self, requires_grad=False): + super(Vgg19, self).__init__() + vgg_pretrained_features = models.vgg19(pretrained=True).features + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + self.slice4 = torch.nn.Sequential() + self.slice5 = torch.nn.Sequential() + for x in range(2): + self.slice1.add_module(str(x), vgg_pretrained_features[x]) + for x in range(2, 7): + self.slice2.add_module(str(x), vgg_pretrained_features[x]) + for x in range(7, 12): + self.slice3.add_module(str(x), vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice4.add_module(str(x), vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice5.add_module(str(x), vgg_pretrained_features[x]) + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + h_relu1 = self.slice1(X) + h_relu2 = self.slice2(h_relu1) + h_relu3 = self.slice3(h_relu2) + h_relu4 = self.slice4(h_relu3) + h_relu5 = self.slice5(h_relu4) + out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] + return out diff --git a/models/runmodel.py b/models/runmodel.py index 7daa092..b41e1cd 100755 --- a/models/runmodel.py +++ b/models/runmodel.py @@ -5,7 +5,6 @@ from util import mosaic from util import data import torch - def run_unet(img,net,size = 128,use_gpu = True): img=impro.image2folat(img,3) img=img.reshape(1,3,size,size) @@ -25,9 +24,12 @@ def run_unet_rectim(img,net,size = 128,use_gpu = True): mask = impro.mergeimage(mask1,mask2,img) return mask -def run_pix2pix(img,net,size = 128,use_gpu = True): - img = impro.resize(img,size) - img = data.im2tensor(img,use_gpu=use_gpu) +def run_pix2pix(img,net,opt): + if opt.netG == 'HD': + img = impro.resize(img,512) + else: + img = impro.resize(img,128) + img = data.im2tensor(img,use_gpu=opt.use_gpu) img_fake = net(img) img_fake = data.tensor2im(img_fake) return img_fake @@ -38,7 +40,6 @@ def get_ROI_position(img,net,opt): x,y,halfsize,area = impro.boundingSquare(mask, 1) return mask,x,y,area - def get_mosaic_position(img_origin,net_mosaic_pos,opt): mask = run_unet_rectim(img_origin,net_mosaic_pos,use_gpu = opt.use_gpu) mask = impro.mask_threshold(mask,10,128) diff --git a/models/unet_parts.py b/models/unet_parts.py index bd5ffee..80165e5 100755 --- a/models/unet_parts.py +++ b/models/unet_parts.py @@ -54,7 +54,7 @@ class Upsample(nn.Module): def forward(self, x): return F.interpolate(x, scale_factor=self.scale_factor,mode='bilinear', align_corners=True) -F.interpolate + class up(nn.Module): def __init__(self, in_ch, out_ch, bilinear=True): super(up, self).__init__() diff --git a/options.py b/options.py index 1ae126d..d13928e 100644 --- a/options.py +++ b/options.py @@ -23,8 +23,8 @@ class Options(): self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255') self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin') - #AddMosaic - self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks') + #CleanMosaic + self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks | HD') self.parser.add_argument('--mosaic_position_model_path', type=str, default='auto',help='name of model use to find mosaic position') self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather,but run faster') self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video') @@ -36,11 +36,16 @@ class Options(): self.initialize() self.opt = self.parser.parse_args() - if self.opt.netG == 'auto': + if self.opt.netG == 'auto' and self.opt.mode =='clean': if 'unet_128' in self.opt.model_path: self.opt.netG = 'unet_128' elif 'resnet_9blocks' in self.opt.model_path: self.opt.netG = 'resnet_9blocks' + elif 'HD' in self.opt.model_path: + self.opt.netG = 'HD' + else: + print('Type of Generator error!') + if self.opt.mosaic_position_model_path == 'auto': _path = os.path.join(os.path.split(self.opt.model_path)[0],'mosaic_position.pth') diff --git a/util/data.py b/util/data.py index 7c563b9..1978c96 100755 --- a/util/data.py +++ b/util/data.py @@ -20,11 +20,12 @@ def tensor2im(image_tensor, imtype=np.uint8, rgb2bgr = True): def im2tensor(image_numpy, imtype=np.uint8, bgr2rgb = True, reshape = True, use_gpu = True): + h, w = image_numpy.shape[:2] if bgr2rgb: image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy) image_tensor = transform(image_numpy) if reshape: - image_tensor=image_tensor.reshape(1,3,128,128) + image_tensor=image_tensor.reshape(1,3,h,w) if use_gpu: image_tensor = image_tensor.cuda() return image_tensor \ No newline at end of file diff --git a/util/mosaic.py b/util/mosaic.py index 72ec23b..99fbf92 100755 --- a/util/mosaic.py +++ b/util/mosaic.py @@ -55,12 +55,16 @@ def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg'): return img_mosaic -def addmosaic_autosize(img,mask,model): +def addmosaic_autosize(img,mask,model,area_type = 'normal'): h,w = img.shape[:2] mask = cv2.resize(mask,(w,h)) alpha = np.min((w,h))/512 try: - area = mask_area(mask) + if area_type == 'normal': + area = mask_area(mask) + elif area_type == 'bounding': + w,h = cv2.boundingRect(mask)[2:] + area = w*h except: area = 0 area = area/(alpha*alpha) @@ -76,19 +80,23 @@ def addmosaic_autosize(img,mask,model): pass return img_mosaic -def addmosaic_random(img,mask): +def addmosaic_random(img,mask,area_type = 'normal'): # img = resize(img,512) h,w = img.shape[:2] mask = cv2.resize(mask,(w,h)) alpha = np.min((w,h))/512 #area_avg=5925*4 try: - area = mask_area(mask) + if area_type == 'normal': + area = mask_area(mask) + elif area_type == 'bounding': + w,h = cv2.boundingRect(mask)[2:] + area = w*h except: area = 0 area = area/(alpha*alpha) if area>50000: - img_mosaic = random_mod(img,mask,alpha*random.uniform(16,28)) + img_mosaic = random_mod(img,mask,alpha*random.uniform(16,30)) elif 20000