Upload 5 files

gan.py

@@ -0,0 +1,240 @@
+from PIL import Image
+import numpy as np
+import cv2
+import torchvision.transforms as transforms
+import torch
+import io
+import os
+import functools
+
+class DataLoader():
+
+	def __init__(self, opt, cv_img):
+		super(DataLoader, self).__init__()
+
+		self.dataset = Dataset()
+		self.dataset.initialize(opt, cv_img)
+
+		self.dataloader = torch.utils.data.DataLoader(
+			self.dataset,
+			batch_size=opt.batchSize,
+			shuffle=not opt.serial_batches,
+			num_workers=int(opt.nThreads))
+
+	def load_data(self):
+		return self.dataloader
+
+	def __len__(self):
+		return 1
+
+class Dataset(torch.utils.data.Dataset):
+	def __init__(self):
+		super(Dataset, self).__init__()
+
+	def initialize(self, opt, cv_img):
+		self.opt = opt
+		self.root = opt.dataroot
+
+		self.A = Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+		self.dataset_size = 1
+	
+	def __getitem__(self, index):        
+
+		transform_A = get_transform(self.opt)
+		A_tensor = transform_A(self.A.convert('RGB'))
+
+		B_tensor = inst_tensor = feat_tensor = 0
+
+		input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor, 
+					  'feat': feat_tensor, 'path': ""}
+
+		return input_dict
+
+	def __len__(self):
+		return 1    
+
+class DeepModel(torch.nn.Module):
+
+	def initialize(self, opt):
+
+		torch.cuda.empty_cache()
+
+		self.opt = opt
+
+		self.gpu_ids = [] #FIX CPU
+
+		self.netG = self.__define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, 
+									  opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, 
+									  opt.n_blocks_local, opt.norm, self.gpu_ids)        
+
+		# load networks
+		self.__load_network(self.netG)
+
+		
+	
+	def inference(self, label, inst):
+		
+		# Encode Inputs        
+		input_label, inst_map, _, _ = self.__encode_input(label, inst, infer=True)
+
+		# Fake Generation
+		input_concat = input_label        
+		
+		with torch.no_grad():
+			fake_image = self.netG.forward(input_concat)
+
+		return fake_image
+	
+	# helper loading function that can be used by subclasses
+	def __load_network(self, network):
+
+		save_path = os.path.join(self.opt.checkpoints_dir)
+
+		network.load_state_dict(torch.load(save_path))
+
+	def __encode_input(self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False):             
+		if (len(self.gpu_ids) > 0): 
+			input_label = label_map.data.cuda() #GPU
+		else: 
+			input_label = label_map.data #CPU
+			
+		return input_label, inst_map, real_image, feat_map
+
+	def __weights_init(self, 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 __define_G(self, 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 = self.__get_norm_layer(norm_type=norm)         
+		netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)
+		
+		if len(gpu_ids) > 0:
+			netG.cuda(gpu_ids[0])
+		netG.apply(self.__weights_init)
+		return netG
+
+	def __get_norm_layer(self, norm_type='instance'):
+		norm_layer = functools.partial(torch.nn.InstanceNorm2d, affine=False)
+		return norm_layer
+
+##############################################################################
+# Generator
+##############################################################################
+class GlobalGenerator(torch.nn.Module):
+	def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=torch.nn.BatchNorm2d, 
+				 padding_type='reflect'):
+		assert(n_blocks >= 0)
+		super(GlobalGenerator, self).__init__()        
+		activation = torch.nn.ReLU(True)        
+
+		model = [torch.nn.ReflectionPad2d(3), torch.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 += [torch.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 += [torch.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 += [torch.nn.ReflectionPad2d(3), torch.nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), torch.nn.Tanh()]        
+		self.model = torch.nn.Sequential(*model)
+			
+	def forward(self, input):
+		return self.model(input)             
+		
+# Define a resnet block
+class ResnetBlock(torch.nn.Module):
+	def __init__(self, dim, padding_type, norm_layer, activation=torch.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 += [torch.nn.ReflectionPad2d(1)]
+		elif padding_type == 'replicate':
+			conv_block += [torch.nn.ReplicationPad2d(1)]
+		elif padding_type == 'zero':
+			p = 1
+		else:
+			raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+		conv_block += [torch.nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+					   norm_layer(dim),
+					   activation]
+		if use_dropout:
+			conv_block += [torch.nn.Dropout(0.5)]
+
+		p = 0
+		if padding_type == 'reflect':
+			conv_block += [torch.nn.ReflectionPad2d(1)]
+		elif padding_type == 'replicate':
+			conv_block += [torch.nn.ReplicationPad2d(1)]
+		elif padding_type == 'zero':
+			p = 1
+		else:
+			raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+		conv_block += [torch.nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+					   norm_layer(dim)]
+
+		return torch.nn.Sequential(*conv_block)
+
+	def forward(self, x):
+		out = x + self.conv_block(x)
+		return out
+
+# Data utils:
+def get_transform(opt, method=Image.BICUBIC, normalize=True):
+	transform_list = []
+
+	base = float(2 ** opt.n_downsample_global)
+	if opt.netG == 'local':
+		base *= (2 ** opt.n_local_enhancers)
+	transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+
+	transform_list += [transforms.ToTensor()]
+
+	if normalize:
+		transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
+												(0.5, 0.5, 0.5))]
+	return transforms.Compose(transform_list)
+
+def __make_power_2(img, base, method=Image.BICUBIC):
+	ow, oh = img.size        
+	h = int(round(oh / base) * base)
+	w = int(round(ow / base) * base)
+	if (h == oh) and (w == ow):
+		return img
+	return img.resize((w, h), method)
+
+# Converts a Tensor into a Numpy array
+# |imtype|: the desired type of the converted numpy array
+def tensor2im(image_tensor, imtype=np.uint8, normalize=True):
+	if isinstance(image_tensor, list):
+		image_numpy = []
+		for i in range(len(image_tensor)):
+			image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
+		return image_numpy
+	image_numpy = image_tensor.cpu().float().numpy()
+	if normalize:
+		image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+	else:
+		image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0      
+	image_numpy = np.clip(image_numpy, 0, 255)
+	if image_numpy.shape[2] == 1 or image_numpy.shape[2] > 3:        
+		image_numpy = image_numpy[:,:,0]
+	return image_numpy.astype(imtype)

gggg.py

@@ -0,0 +1,78 @@
+from run import process
+import time
+import subprocess
+import os
+import argparse
+import cv2
+import sys
+from PIL import Image
+import torch
+import gradio as gr
+
+
+TESTdevice = "cpu"
+
+index = 1
+
+
+"""
+main.py
+
+ How to run:
+ python main.py
+
+"""
+
+
+def mainTest(inputpath, outpath):
+    watermark = deep_nude_process(inputpath)
+    cv2.imwrite(outpath, watermark)
+    return watermark
+    #
+
+
+def deep_nude_process(item):
+    # print('Processing {}'.format(item))
+    # dress = cv2.imread(item)
+    dress = (item)
+    h = dress.shape[0]
+    w = dress.shape[1]
+    dress = cv2.resize(dress, (512, 512), interpolation=cv2.INTER_CUBIC)
+    watermark = process(dress)
+    watermark = cv2.resize(watermark, (w, h), interpolation=cv2.INTER_CUBIC)
+    return watermark
+
+
+def inference(img):
+    global index
+    # inputpath = "input/" + str(index) + ".jpg"
+    outputpath = "out_" + str(index) + ".jpg"
+    # cv2.imwrite(inputpath, img)
+    index += 1
+    print(time.strftime("START!!!!!!!!! %Y-%m-%d %H:%M:%S", time.localtime()))
+    output = mainTest(img, outputpath)
+    print(time.strftime("FINISH!!!!!!!!! %Y-%m-%d %H:%M:%S", time.localtime()))
+    return output
+
+
+title = "AI脱衣"
+description = "传入人物照片,类似最下方测试图的那种,将制作脱衣图,一张图至少等40秒,别传私人照片,禁止传真人照片"
+
+examples = [
+    ['input.png', '测试图'],
+]
+
+
+web = gr.Interface(inference,
+                   inputs="image",
+                   outputs="image",
+                   title=title,
+                   description=description,
+                   examples=examples,
+                   )
+
+if __name__ == '__main__':
+    web.launch(
+        share=True,
+        enable_queue=True
+    )

main.py

@@ -0,0 +1,49 @@
+# coding=utf-8
+import sys
+import cv2
+import argparse
+import os
+import sys
+import subprocess
+import time
+
+from run import process
+
+"""
+main.py
+
+ How to run:
+ python main.py
+
+"""
+
+
+def main(inputpath, outpath, show):
+	if isinstance(inputpath, list):
+		for item in inputpath:
+			watermark = deep_nude_process(item)
+			cv2.imwrite("output_"+item, watermark)
+	else:
+		watermark = deep_nude_process(inputpath)
+		cv2.imwrite(outpath, watermark)
+
+def deep_nude_process(item):
+    print('Processing {}'.format(item))
+    dress = cv2.imread(item)
+    h = dress.shape[0]
+    w = dress.shape[1]
+    dress = cv2.resize(dress, (512,512), interpolation=cv2.INTER_CUBIC)
+    watermark = process(dress)
+    watermark = cv2.resize(watermark, (w,h), interpolation=cv2.INTER_CUBIC)
+    return watermark
+
+if __name__ == '__main__':
+	parser = argparse.ArgumentParser(description="simple deep nude script tool")
+	parser.add_argument("-i", "--input", action="store", nargs = "*", default="input.png", help = "Use to enter input one or more files's name")
+	parser.add_argument("-o", "--output", action="store", default="output.png", help = "Use to enter output file name")
+	parser.add_argument("-s", "--show", action="store", default="false", help = "Use to automatically display or not display generated images")
+	inputpath, outputpath, show = parser.parse_args().input, parser.parse_args().output, parser.parse_args().show
+    
+	print (time.strftime("START!!!!!!!!! %Y-%m-%d %H:%M:%S", time.localtime()))
+	main(inputpath, outputpath, show)
+	print (time.strftime("FINISH!!!!!!!!! %Y-%m-%d %H:%M:%S", time.localtime()))

run.py

@@ -0,0 +1,150 @@
+import cv2
+
+#Import Neural Network Model
+from gan import DataLoader, DeepModel, tensor2im
+
+#OpenCv Transform:
+from opencv_transform.mask_to_maskref import create_maskref
+from opencv_transform.maskdet_to_maskfin import create_maskfin
+from opencv_transform.dress_to_correct import create_correct
+from opencv_transform.nude_to_watermark import create_watermark
+
+"""
+run.py
+
+This script manage the entire transormation.
+
+Transformation happens in 6 phases:
+	0: dress -> correct [opencv] dress_to_correct
+	1: correct -> mask:  [GAN] correct_to_mask
+	2: mask -> maskref [opencv] mask_to_maskref
+	3: maskref -> maskdet [GAN] maskref_to_maskdet
+	4: maskdet -> maskfin [opencv] maskdet_to_maskfin
+	5: maskfin -> nude [GAN] maskfin_to_nude
+	6: nude -> watermark [opencv] nude_to_watermark
+
+"""
+
+phases = ["dress_to_correct", "correct_to_mask", "mask_to_maskref", "maskref_to_maskdet", "maskdet_to_maskfin", "maskfin_to_nude", "nude_to_watermark"]
+
+class Options():
+
+	#Init options with default values
+	def __init__(self):
+	
+		# experiment specifics
+		self.norm = 'batch' #instance normalization or batch normalization
+		self.use_dropout = False #use dropout for the generator
+		self.data_type = 32 #Supported data type i.e. 8, 16, 32 bit
+
+		# input/output sizes       
+		self.batchSize = 1 #input batch size
+		self.input_nc = 3 # of input image channels
+		self.output_nc = 3 # of output image channels
+
+		# for setting inputs
+		self.serial_batches = True #if true, takes images in order to make batches, otherwise takes them randomly
+		self.nThreads = 1 ## threads for loading data (???)
+		self.max_dataset_size = 1 #Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.
+		
+		# for generator
+		self.netG = 'global' #selects model to use for netG
+		self.ngf = 64 ## of gen filters in first conv layer
+		self.n_downsample_global = 4 #number of downsampling layers in netG
+		self.n_blocks_global = 9 #number of residual blocks in the global generator network
+		self.n_blocks_local = 0 #number of residual blocks in the local enhancer network
+		self.n_local_enhancers = 0 #number of local enhancers to use
+		self.niter_fix_global = 0 #number of epochs that we only train the outmost local enhancer
+
+		#Phase specific options
+		self.checkpoints_dir = ""
+		self.dataroot = ""
+
+	#Changes options accordlying to actual phase
+	def updateOptions(self, phase):
+
+		if phase == "correct_to_mask":
+			self.checkpoints_dir = "checkpoints/cm.lib"
+
+		elif phase == "maskref_to_maskdet":
+			self.checkpoints_dir = "checkpoints/mm.lib"
+
+		elif phase == "maskfin_to_nude":
+			self.checkpoints_dir = "checkpoints/mn.lib"
+
+# process(cv_img, mode)
+# return:
+# 	watermark image
+def process(cv_img):
+
+	#InMemory cv2 images:
+	dress = cv_img
+	correct = None
+	mask = None
+	maskref = None
+	maskfin = None
+	maskdet = None
+	nude = None
+	watermark = None
+
+	for index, phase in enumerate(phases):
+
+		print("Executing phase: " + phase) 
+			
+		#GAN phases:
+		if (phase == "correct_to_mask") or (phase == "maskref_to_maskdet") or (phase == "maskfin_to_nude"):
+
+			#Load global option
+			opt = Options()
+
+			#Load custom phase options:
+			opt.updateOptions(phase)
+
+			#Load Data
+			if (phase == "correct_to_mask"):
+				data_loader = DataLoader(opt, correct)
+			elif (phase == "maskref_to_maskdet"):
+				data_loader = DataLoader(opt, maskref)
+			elif (phase == "maskfin_to_nude"):
+				data_loader = DataLoader(opt, maskfin)
+			
+			dataset = data_loader.load_data()
+			
+			#Create Model
+			model = DeepModel()
+			model.initialize(opt)
+
+			#Run for every image:
+			for i, data in enumerate(dataset):
+
+				generated = model.inference(data['label'], data['inst'])
+
+				im = tensor2im(generated.data[0])
+
+				#Save Data
+				if (phase == "correct_to_mask"):
+					mask = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
+
+				elif (phase == "maskref_to_maskdet"):
+					maskdet = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
+
+				elif (phase == "maskfin_to_nude"):
+					nude = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
+
+		#Correcting:
+		elif (phase == 'dress_to_correct'):
+			correct = create_correct(dress)
+
+		#mask_ref phase (opencv)
+		elif (phase == "mask_to_maskref"):
+			maskref = create_maskref(mask, correct)
+
+		#mask_fin phase (opencv)
+		elif (phase == "maskdet_to_maskfin"):
+			maskfin = create_maskfin(maskref, maskdet)
+
+		#nude_to_watermark phase (opencv)
+		elif (phase == "nude_to_watermark"):
+			watermark = create_watermark(nude)
+
+	return watermark