Added square/stories format

app.py

@@ -32,15 +32,25 @@ def get_mask(img_in):
     if response.status_code == 200:
         result = response.json()
 
-        mask = Image.open(io.BytesIO(base64.b64decode(
-            result['mask']))).resize(img_in.size)
+        img_square = Image.open(io.BytesIO(base64.b64decode(
+            result['square'])))
+        img_stories = Image.open(io.BytesIO(base64.b64decode(
+            result['stories'])))
 
-        return (mask, result['emotion'])
+        return (img_square, img_stories, result['emotion'])
     else:
         print(f"ERROR: {response.json()['detail']}")
         raise gr.Error(response.json()['detail'])
 
 
+header = r"""
+<center>
+<h1>ARKA Selfie App</h1>
+<p>Send this app a selfie to add a custom background and elements to it.</p>
+<p>It also detects and displays human emotions and works anywhere.</p>
+</center>
+"""
+
 footer = r"""
 <center>
 <b>
@@ -50,18 +60,19 @@ DEMO FOR BG REMOVAL
 """
 
 with gr.Blocks(title="ARKA + TikTok Demo") as app:
-    gr.HTML("<center><h1>ARKA Remove Background</h1></center>")
+    gr.HTML(header)
+    with gr.Row():
+        run_btn = gr.Button(variant="primary", interactive=True, label="Run")
     with gr.Row():
         with gr.Column():
             input_img = gr.Image(type="pil", label="Input image")
-            # new_img = gr.Image(type="pil", label="Custom background")
         with gr.Column():
-            output_img = gr.Image(type="pil", label="result")
+            with gr.Row():
+                output_img = gr.Image(type="pil", label="result_square")
+                output_img1 = gr.Image(type="pil", label="result_stories")
             emotion = gr.Label(label="emotion")
-    with gr.Row():
-        run_btn = gr.Button(variant="primary")
 
-    run_btn.click(get_mask, [input_img], [output_img, emotion])
+    run_btn.click(get_mask, [input_img], [output_img, output_img1, emotion])
 
     with gr.Row():
         gr.HTML(footer)