Update app.py

app.py

@@ -49,7 +49,7 @@ def display_gif_pad(file_name, save_name):
         image_filename = file_name + frame_name + '.png'
         image = imageio.imread(image_filename)
         image = image[:, :, :3]
-        image_pad = cv2.copyMakeBorder(image, 0, 0, 80, 80, cv2.BORDER_CONSTANT, value=0)
+        image_pad = cv2.copyMakeBorder(image, 0, 0, 100, 100, cv2.BORDER_CONSTANT, value=0)
         images.append(image_pad)
 
     gif_filename = 'avatar_source.gif'
@@ -64,11 +64,13 @@ def display_image(file_name):
     imageio.imwrite('image.png', image)
     
 
-def run(action, body, hair, top, bottom):
+def run(action, hair, top, bottom):
 
+    # == Source Avatar ==
     # body
-    if body == "human": body = '0'
-    elif body == "alien": body = '1'
+    #if body == "human":body = '0'
+    #elif body == "alien": body = '1'
+    body = '0'
     
     # hair
     if hair == "green": hair = '0'
@@ -93,20 +95,31 @@ def run(action, body, hair, top, bottom):
     
     gif = display_gif_pad(file_name_source, 'avatar_source.gif')
     
-    return 'avatar_source.gif'
+    # == Target Avatar ==
+    body_target = '1'
+    hair_target = np.random.choice('1', '3', '5', '6', '9')
+    top_target = np.random.choice('3', '4', '5', '6')
+    bottom_target = np.random.choice('4', '5', '6')
+    
+    file_name_target = './Sprite/frames/domain_2/' + action + '/'
+    file_name_target = file_name_target + 'front' + '_' + str(body_target) + str(bottom_target) + str(top_target) + str(hair_target) + '_'
+    
+    gif = display_gif_pad(file_name_target, 'avatar_target.gif')
+    
+    return 'avatar_source.gif', 'avatar_target.gif'
 
 
 gr.Interface(
     run,
     inputs=[
         gr.Radio(choices=["shoot", "slash", "spellcard", "thrust", "walk"], value="shoot"),
-        gr.Radio(choices=["human", "alien"], value="human"),
         gr.Radio(choices=["green", "yellow", "rose", "red", "wine"], value="green"),
         gr.Radio(choices=["brown", "blue", "white"], value="brown"),
         gr.Radio(choices=["white", "golden", "red", "silver"], value="white"),
     ],
     outputs=[
-        gr.components.Image(type="file", label="Avatar (Source)"),
+        gr.components.Image(type="file", label="Source Avatar (Costumed by You)"),
+        gr.components.Image(type="file", label="Target Avatar (Randomly Chosen)"),
     ],
     live=True,
     title="TransferVAE for Unsupervised Video Domain Adaptation",