##################################################### # AST Composite Server Double Two # By Guillaume Descoteaux-Isabelle, 20021 # # This server compose two Adaptive Style Transfer model (output of the first pass serve as input to the second using the same model) ######################################################## #v1-dev #Receive the 2 res from arguments in the request... import os import numpy as np import tensorflow as tf import cv2 from module import encoder, decoder from glob import glob import runway from runway.data_types import number, text #from utils import * import scipy from datetime import datetime import time import re SRV_TYPE="s1" #set env var RW_ if not already set if not os.getenv('RW_PORT'): os.environ["RW_PORT"] = "7860" if not os.getenv('RW_DEBUG'): os.environ["RW_DEBUG"] = "0" if not os.getenv('RW_HOST'): os.environ["RW_HOST"] = "0.0.0.0" #RW_MODEL_OPTIONS if not os.getenv('RW_MODEL_OPTIONS'): os.environ["RW_MODEL_OPTIONS"]='{"styleCheckpoint":"/data/styleCheckpoint"}' # Determining the size of the passes pass1_image_size = 1328 if not os.getenv('PASS1IMAGESIZE'): print("PASS1IMAGESIZE env var non existent;using default:" + str(pass1_image_size)) else: pass1_image_size = os.getenv('PASS1IMAGESIZE', 1328) print("PASS1IMAGESIZE value:" + str(pass1_image_size)) # Determining the size of the passes autoabc = 1 if not os.getenv('AUTOABC'): print("AUTOABC env var non existent;using default:") print(autoabc) abcdefault = 1 print("NOTE----> when running docker, set AUTOABC variable") print(" docker run ... -e AUTOABC=1 #enabled, 0 to disabled (default)") else: autoabc = os.getenv('AUTOABC',1) print("AUTOABC value:") print(autoabc) abcdefault = autoabc #pass2_image_size = 1024 #if not os.getenv('PASS2IMAGESIZE'): # print("PASS2IMAGESIZE env var non existent;using default:" + pass2_image_size) #else: # pass2_image_size = os.getenv('PASS2IMAGESIZE') # print("PASS2IMAGESIZE value:" + pass2_image_size) # pass3_image_size = 2048 # if not os.getenv('PASS3IMAGESIZE'): # print("PASS3IMAGESIZE env var non existent;using default:" + pass3_image_size) # else: # pass3_image_size = os.getenv('PASS3IMAGESIZE') # print("PASS3IMAGESIZE value:" + pass3_image_size) ########################################## ## MODELS #model name for sending it in the response model1name = "UNNAMED" if not os.getenv('MODEL1NAME'): print("MODEL1NAME env var non existent;using default:" + model1name) else: model1name = os.getenv('MODEL1NAME', "UNNAMED") print("MODEL1NAME value:" + model1name) # #m2 # model2name = "UNNAMED" # if not os.getenv('MODEL2NAME'): print("MODEL2NAME env var non existent;using default:" + model2name) # else: # model2name = os.getenv('MODEL2NAME') # print("MODEL2NAME value:" + model2name) # #m3 # model3name = "UNNAMED" # if not os.getenv('MODEL3NAME'): print("MODEL3NAME env var non existent;using default:" + model3name) # else: # model3name = os.getenv('MODEL3NAME') # print("MODEL3NAME value:" + model3name) ####################################################### def get_model_simplified_name_from_dirname(dirname): result_simple_name = dirname.replace("model_","").replace("_864x","").replace("_864","").replace("_new","").replace("-864","") print(" result_simple_name:" + result_simple_name) return result_simple_name def get_padded_checkpoint_no_from_filename(checkpoint_filename): match = re.search(r'ckpt-(\d+)', checkpoint_filename) if match: number = int(match.group(1)) checkpoint_number = round(number/1000,0) print(checkpoint_number) padded_checkpoint_number = str(str(checkpoint_number).zfill(3)) return padded_checkpoint_number.replace('.0','') found_model='none' found_model_checkpoint='0' ######################################################### # SETUP runway_files = runway.file(is_directory=True) @runway.setup(options={'styleCheckpoint': runway_files}) def setup(opts): global found_model,found_model_checkpoint sess = tf.Session() # sess2 = tf.Session() # sess3 = tf.Session() init_op = tf.global_variables_initializer() # init_op2 = tf.global_variables_initializer() # init_op3 = tf.global_variables_initializer() sess.run(init_op) # sess2.run(init_op2) # sess3.run(init_op3) with tf.name_scope('placeholder'): input_photo = tf.placeholder(dtype=tf.float32, shape=[1, None, None, 3], name='photo') input_photo_features = encoder(image=input_photo, options={'gf_dim': 32}, reuse=False) output_photo = decoder(features=input_photo_features, options={'gf_dim': 32}, reuse=False) saver = tf.train.Saver() # saver2 = tf.train.Saver() # saver3 = tf.train.Saver() print("-------------====PATH---------------------->>>>--") path_default = '/data/styleCheckpoint' print("opts:") print(opts) print("----------------------------------------") if opts is None: print("ERROR:opts is None") path = path_default try: path = opts['styleCheckpoint'] except: opts= {'styleCheckpoint': u'/data/styleCheckpoint'} path = opts['styleCheckpoint'] if not os.path.exists(path): print("ERROR:Path does not exist:" + path) path = path_default print(path) print("----------------PATH=======---------------<<<<--") #Getting the model name model_name = [p for p in os.listdir(path) if os.path.isdir(os.path.join(path, p))][0] if not os.getenv('MODELNAME'): dtprint("CONFIG::MODELNAME env var non existent;using default:" + model_name) else: model_name = os.getenv('MODELNAME') # #Getting the model2 name # model2_name = [p for p in os.listdir(path) if os.path.isdir(os.path.join(path, p))][1] # if not os.getenv('MODEL2NAME'): # dtprint("CONFIG::MODEL2NAME env var non existent;using default:" + model2_name) # else: # model2_name = os.getenv('MODEL2NAME') ##Getting the model3 name # model3_name = [p for p in os.listdir(path) if os.path.isdir(os.path.join(path, p))][2] # if not os.getenv('MODEL3NAME'): # dtprint("CONFIG::MODEL3NAME env var non existent;using default:" + model3_name) # else: # model3_name = os.getenv('MODEL3NAME') checkpoint_dir = os.path.join(path, model_name, 'checkpoint_long') #checkpoint2_dir = os.path.join(path, model2_name, 'checkpoint_long') # checkpoint3_dir = os.path.join(path, model3_name, 'checkpoint_long') print("-----------------------------------------") print("modelname is : " + model_name) found_model=get_model_simplified_name_from_dirname(model_name) #print("model2name is : " + model2_name) # print("model3name is : " + model3_name) print("checkpoint_dir is : " + checkpoint_dir) #print("checkpoint2_dir is : " + checkpoint2_dir) # print("checkpoint3_dir is : " + checkpoint3_dir) print("-----------------------------------------") ckpt = tf.train.get_checkpoint_state(checkpoint_dir) #ckpt2 = tf.train.get_checkpoint_state(checkpoint2_dir) # ckpt3 = tf.train.get_checkpoint_state(checkpoint3_dir) ckpt_name = os.path.basename(ckpt.model_checkpoint_path) found_model_checkpoint= get_padded_checkpoint_no_from_filename(ckpt_name) #ckpt2_name = os.path.basename(ckpt2.model_checkpoint_path) # ckpt3_name = os.path.basename(ckpt3.model_checkpoint_path) saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name)) #saver2.restore(sess2, os.path.join(checkpoint2_dir, ckpt2_name)) # saver3.restore(sess3, os.path.join(checkpoint3_dir, ckpt3_name)) m1 = dict(sess=sess, input_photo=input_photo, output_photo=output_photo) #m2 = dict(sess=sess2, input_photo=input_photo, output_photo=output_photo) # m3 = dict(sess=sess3, input_photo=input_photo, output_photo=output_photo) models = type('', (), {})() models.m1 = m1 #models.m2 = m2 # models.m3 = m3 return models def make_target_output_filename( mname,checkpoint, fn='',res1=0,abc=0, ext='.jpg',svrtype="s1", modelid='', suffix='', xtra_model_id='',verbose=False): fn_base=fn.replace(ext,"") fn_base=fn_base.replace(".jpg","") fn_base=fn_base.replace(".jpeg","") fn_base=fn_base.replace(".JPG","") fn_base=fn_base.replace(".JPEG","") fn_base=fn_base.replace(".png","") fn_base=fn_base.replace(".PNG","") #pad res1 and res2 to 4 digits res1_pad=str(res1).zfill(4) abc_pad=str(abc).zfill(2) if res1_pad=="0000": res1_pad="" #pad checkpoint to 3 digits checkpoint=checkpoint.zfill(3) if fn_base=="none": fn_base="" if '/' in fn_base: fn_base=fn_base.split('/')[-1] # Print out all input info: if verbose : print("-----------------------------") print("fn_base: ",fn_base) print("mname: ",mname) print("suffix: ",suffix) print("res1: ",res1_pad) print("abc: ",abc_pad) print("ext: ",ext) print("svrtype: ",svrtype) print("modelid: ",modelid) print("xtra_model_id: ",xtra_model_id) print("checkpoint: ",checkpoint) print("fn: ",fn) mtag = "{}__{}__{}x{}__{}__{}k".format(mname,suffix,res1_pad,abc_pad, svrtype, checkpoint).replace("_0x" + str(abc_pad), "") if verbose: print(mtag) target_output = "{}__{}__{}{}{}".format(fn_base, modelid, mtag, xtra_model_id, ext).replace("_"+str(abc_pad)+"x"+str(abc_pad)+"_","").replace("_0x0_", "").replace("_0_", "").replace("_-", "_").replace("____", "__").replace("___", "__").replace("___", "__").replace("..",".").replace("model_","").replace("_x"+str(abc_pad)+"_","").replace("gia-ds-","") target_output = replace_values_from_csv(target_output) return target_output def replace_values_from_csv(target_output): # Implement the logic to replace values from CSV #load replacer.csv and replace the values (src,dst) src_dest_file = 'replacer.csv' if os.path.exists(src_dest_file): with open(src_dest_file, 'r') as file: lines = file.readlines() for line in lines: src, dst = line.split(',') target_output = target_output.replace(src, dst) return target_output.replace("\n", "").replace("\r", "").replace(" ", "_") def _make_meta_as_json(x1=0,c1=0,inp=None,result_dict=None): global found_model,found_model_checkpoint fn='none' if inp['fn'] != 'none': fn=inp['fn'] ext='.jpg' if inp['ext'] != '.jpg': ext=inp['ext'] filename=make_target_output_filename(found_model,found_model_checkpoint,fn,x1,c1,ext,SRV_TYPE) if result_dict is None: json_return = { "model": str(found_model), "checkpoint": str(found_model_checkpoint), "filename": str(filename) } return json_return else: #support adding to the existing dict the data directly result_dict['model']=str(found_model) result_dict['checkpoint']=str(found_model_checkpoint) result_dict['filename']=str(filename) return result_dict meta_inputs={'meta':text} meta_outputs={'model':text,'filename':text,'checkpoint':text} @runway.command('meta2', inputs=meta_inputs, outputs=meta_outputs) def get_geta(models, inp): global found_model,found_model_checkpoint json_return = _make_meta_as_json() # "files": "nothing yet" print(json_return) return json_return @runway.command('meta', inputs=meta_inputs, outputs=meta_outputs) def get_geta(models, inp): global found_model,found_model_checkpoint json_return = _make_meta_as_json(inp) # "files": "nothing yet" print(json_return) return json_return #@STCGoal add number or text to specify resolution of the three pass inputs={'contentImage': runway.image,'x1':number(default=1024,min=24,max=18000),'c1':number(default=0,min=-99,max=99),'fn':text(default='none'),'ext':text(default='.jpg')} outputs={'stylizedImage': runway.image,'totaltime':number,'x1': number,'c1':number,'model1name':text,'checkpoint':text,'filename':text,'model':text} @runway.command('stylize', inputs=inputs, outputs=outputs) def stylize(models, inp): global found_model,found_model_checkpoint,model1name start = time.time() model = models.m1 #model2 = models.m2 # model3 = models.m3 #Getting our names back (even though I think we dont need) #@STCIssue BUGGED # m1name=models.m1.name # m2name=models.m2.name # m3name=models.m3.name #get size from inputs rather than env x1 = int(inp['x1']) c1 = int(inp['c1']) # img = inp['contentImage'] img = np.array(img) img = img / 127.5 - 1. #@a Pass 1 RESIZE to 1368px the smaller side image_size=pass1_image_size image_size=x1 img_shape = img.shape[:2] alpha = float(image_size) / float(min(img_shape)) #dtprint ("DEBUG::content.imgshape:" + str(tuple(img_shape)) + ", alpha:" + str(alpha)) try: img = scipy.misc.imresize(img, size=alpha) except: pass img = np.expand_dims(img, axis=0) #@a INFERENCE PASS 1 dtprint("INFO:Pass1 inference starting") img = model['sess'].run(model['output_photo'], feed_dict={model['input_photo']: img}) # img = (img + 1.) * 127.5 img = img.astype('uint8') img = img[0] #dtprint("INFO:Upresing Pass1 for Pass 2 (STARTING) ") #@a Pass 2 RESIZE to 1024px the smaller side #image_size=pass2_image_size #image_size=x2 #img_shape = img.shape[:2] #alpha = float(image_size) / float(min(img_shape)) #dtprint ("DEBUG::pass1.imgshape:" + str(tuple(img_shape)) + ", alpha:" + str(alpha)) #img = scipy.misc.imresize(img, size=alpha) #dtprint("INFO:Upresing Pass1 (DONE) ") #Iteration 2 #img = np.array(img) #img = img / 127.5 - 1. #img = np.expand_dims(img, axis=0) #@a INFERENCE PASS 2 using the same model #dtprint("INFO:Pass2 inference (STARTING)") #img = model['sess'].run(model['output_photo'], feed_dict={model['input_photo']: img}) #dtprint("INFO:Pass2 inference (DONE)") #img = (img + 1.) * 127.5 #img = img.astype('uint8') #img = img[0] # #pass3 # #@a Pass 3 RESIZE to 2048px the smaller side # image_size=pass3_image_size # image_size=x3 # img_shape = img.shape[:2] # alpha = float(image_size) / float(min(img_shape)) # dtprint ("DEBUG::pass2.imgshape:" + str(tuple(img_shape)) + ", alpha:" + str(alpha)) # img = scipy.misc.imresize(img, size=alpha) # dtprint("INFO:Upresing Pass2 (DONE) ") # #Iteration 3 # img = np.array(img) # img = img / 127.5 - 1. # img = np.expand_dims(img, axis=0) # #@a INFERENCE PASS 3 # dtprint("INFO:Pass3 inference (STARTING)") # img = model3['sess'].run(model3['output_photo'], feed_dict={model3['input_photo']: img}) # dtprint("INFO:Pass3 inference (DONE)") # img = (img + 1.) * 127.5 # img = img.astype('uint8') # img = img[0] # #pass3 #dtprint("INFO:Composing done") if c1 != 0 : print('Auto Brightening images...' + str(c1)) img = img, alpha2, beta = automatic_brightness_and_contrast(img,c1) stop = time.time() totaltime = stop - start print("The time of the run:", totaltime) #if model1name UNNAMED, use found_model if model1name == "UNNAMED": model1name=found_model include_meta_directly_in_result=True if include_meta_directly_in_result: result_dict = dict(stylizedImage=img,totaltime=totaltime,x1=x1,model1name=model1name,c1=c1) result_dict = _make_meta_as_json(x1,c1,inp,result_dict) else: meta_data = _make_meta_as_json(x1,c1,inp) result_dict = dict(stylizedImage=img,totaltime=totaltime,x1=x1,model1name=model1name,c1=c1,meta=meta_data) return result_dict def dtprint(msg): dttag=getdttag() print(dttag + "::" + msg ) def getdttag(): # datetime object containing current date and time now = datetime.now() # dd/mm/YY H:M:S # dt_string = now.strftime("%d/%m/%Y %H:%M:%S") return now.strftime("%H:%M:%S") # Automatic brightness and contrast optimization with optional histogram clipping def automatic_brightness_and_contrast(image, clip_hist_percent=25): gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # Calculate grayscale histogram hist = cv2.calcHist([gray],[0],None,[256],[0,256]) hist_size = len(hist) # Calculate cumulative distribution from the histogram accumulator = [] accumulator.append(float(hist[0])) for index in range(1, hist_size): accumulator.append(accumulator[index -1] + float(hist[index])) # Locate points to clip maximum = accumulator[-1] clip_hist_percent *= (maximum/100.0) clip_hist_percent /= 2.0 # Locate left cut minimum_gray = 0 while accumulator[minimum_gray] < clip_hist_percent: minimum_gray += 1 # Locate right cut maximum_gray = hist_size -1 while accumulator[maximum_gray] >= (maximum - clip_hist_percent): maximum_gray -= 1 # Calculate alpha and beta values alpha = 255 / (maximum_gray - minimum_gray) beta = -minimum_gray * alpha ''' # Calculate new histogram with desired range and show histogram new_hist = cv2.calcHist([gray],[0],None,[256],[minimum_gray,maximum_gray]) plt.plot(hist) plt.plot(new_hist) plt.xlim([0,256]) plt.show() ''' auto_result = cv2.convertScaleAbs(image, alpha=alpha, beta=beta) return (auto_result, alpha, beta) if __name__ == '__main__': #print('External Service port is:' +os.environ.get('SPORT')) os.environ["RW_PORT"] = "7860" print("Launched...") runway.run()