Update app.py (#4)

- Update app.py (4713eaf95f412781d5a37c855dccd9c0552640ca)


Co-authored-by: Aditya Deshmukh <AdityA36912355@users.noreply.huggingface.co>

app.py

@@ -115,8 +115,10 @@ def predict(
     seed,
     negative_prompt,
     task,
-    vertical_expansion_ratio,
-    horizontal_expansion_ratio,
+    left_expansion_ratio,
+    right_expansion_ratio,
+    top_expansion_ratio,
+    bottom_expansion_ratio,
 ):
     size1, size2 = input_image["image"].convert("RGB").size
 
@@ -136,40 +138,30 @@ def predict(
     if task == "object-removal":
         prompt = prompt + " empty scene blur"
 
-    if vertical_expansion_ratio != None and horizontal_expansion_ratio != None:
+    if (
+        left_expansion_ratio is not None and right_expansion_ratio is not None
+        and top_expansion_ratio is not None and bottom_expansion_ratio is not None
+    ):
         o_W, o_H = input_image["image"].convert("RGB").size
-        c_W = int(horizontal_expansion_ratio * o_W)
-        c_H = int(vertical_expansion_ratio * o_H)
+        c_W = int((1 + left_expansion_ratio + right_expansion_ratio) * o_W)
+        c_H = int((1 + top_expansion_ratio + bottom_expansion_ratio) * o_H)
 
         expand_img = np.ones((c_H, c_W, 3), dtype=np.uint8) * 127
         original_img = np.array(input_image["image"])
         expand_img[
-            int((c_H - o_H) / 2.0) : int((c_H - o_H) / 2.0) + o_H,
-            int((c_W - o_W) / 2.0) : int((c_W - o_W) / 2.0) + o_W,
-            :,
+            int(top_expansion_ratio * o_H):int(top_expansion_ratio * o_H) + o_H,
+            int(left_expansion_ratio * o_W):int(left_expansion_ratio * o_W) + o_W,
+            :
         ] = original_img
 
         blurry_gap = 10
 
         expand_mask = np.ones((c_H, c_W, 3), dtype=np.uint8) * 255
-        if vertical_expansion_ratio == 1 and horizontal_expansion_ratio != 1:
-            expand_mask[
-                int((c_H - o_H) / 2.0) : int((c_H - o_H) / 2.0) + o_H,
-                int((c_W - o_W) / 2.0) + blurry_gap : int((c_W - o_W) / 2.0) + o_W - blurry_gap,
-                :,
-            ] = 0
-        elif vertical_expansion_ratio != 1 and horizontal_expansion_ratio != 1:
-            expand_mask[
-                int((c_H - o_H) / 2.0) + blurry_gap : int((c_H - o_H) / 2.0) + o_H - blurry_gap,
-                int((c_W - o_W) / 2.0) + blurry_gap : int((c_W - o_W) / 2.0) + o_W - blurry_gap,
-                :,
-            ] = 0
-        elif vertical_expansion_ratio != 1 and horizontal_expansion_ratio == 1:
-            expand_mask[
-                int((c_H - o_H) / 2.0) + blurry_gap : int((c_H - o_H) / 2.0) + o_H - blurry_gap,
-                int((c_W - o_W) / 2.0) : int((c_W - o_W) / 2.0) + o_W,
-                :,
-            ] = 0
+        expand_mask[
+            int(top_expansion_ratio * o_H) + blurry_gap:int(top_expansion_ratio * o_H) + o_H - blurry_gap,
+            int(left_expansion_ratio * o_W) + blurry_gap:int(left_expansion_ratio * o_W) + o_W - blurry_gap,
+            :
+        ] = 0
 
         input_image["image"] = Image.fromarray(expand_img)
         input_image["mask"] = Image.fromarray(expand_mask)
@@ -184,11 +176,16 @@ def predict(
 
     np_inpimg = np.array(input_image["image"])
     np_inmask = np.array(input_image["mask"]) / 255.0
+    if len(np_inmask.shape)==2:
+      np_inmask = np.expand_dims(np_inmask, axis=-1)
+    # return np_inpimg, np_inmask
 
     np_inpimg = np_inpimg * (1 - np_inmask)
 
     input_image["image"] = Image.fromarray(np_inpimg.astype(np.uint8)).convert("RGB")
 
+
+    # return input_image
     set_seed(seed)
     global pipe
     result = pipe(
@@ -245,8 +242,10 @@ def infer(
     scale,
     seed,
     task,
-    vertical_expansion_ratio,
-    horizontal_expansion_ratio,
+    left_expansion_ratio,
+    right_expansion_ratio,
+    top_expansion_ratio,
+    bottom_expansion_ratio,
     outpaint_prompt,
     outpaint_negative_prompt,
     removal_prompt,
@@ -278,8 +277,10 @@ def infer(
             seed,
             negative_prompt,
             task,
-            vertical_expansion_ratio,
-            horizontal_expansion_ratio,
+            left_expansion_ratio,
+            right_expansion_ratio,
+            top_expansion_ratio,
+            bottom_expansion_ratio
         )
     else:
         task = "text-guided"
@@ -371,15 +372,29 @@ with gr.Blocks(css="style.css") as demo:
                 )
                 outpaint_prompt = gr.Textbox(label="Outpainting_prompt")
                 outpaint_negative_prompt = gr.Textbox(label="Outpainting_negative_prompt")
-                horizontal_expansion_ratio = gr.Slider(
-                    label="horizontal expansion ratio",
+                left_expansion_ratio = gr.Slider(
+                    label="left expansion ratio",
+                    minimum=1,
+                    maximum=4,
+                    step=0.05,
+                    value=1,
+                )
+                right_expansion_ratio = gr.Slider(
+                    label="right expansion ratio",
+                    minimum=1,
+                    maximum=4,
+                    step=0.05,
+                    value=1,
+                )
+                top_expansion_ratio = gr.Slider(
+                    label="top expansion ratio",
                     minimum=1,
                     maximum=4,
                     step=0.05,
                     value=1,
                 )
-                vertical_expansion_ratio = gr.Slider(
-                    label="vertical expansion ratio",
+                bottom_expansion_ratio = gr.Slider(
+                    label="bottom expansion ratio",
                     minimum=1,
                     maximum=4,
                     step=0.05,
@@ -433,8 +448,10 @@ with gr.Blocks(css="style.css") as demo:
             scale,
             seed,
             task,
-            vertical_expansion_ratio,
-            horizontal_expansion_ratio,
+            left_expansion_ratio,
+            right_expansion_ratio,
+            top_expansion_ratio,
+            bottom_expansion_ratio,
             outpaint_prompt,
             outpaint_negative_prompt,
             removal_prompt,