Update app.py

app.py

@@ -201,79 +201,86 @@ def main_fun(image_pil, x_units, y_units, text_prompt):
     output = draw_image(image_pil, masks, boxes, alpha=0.4)
 
 
-#     '''Get masked object and background as two separate images'''
-#     mask = np.expand_dims(masks[0], axis=-1)
-#     masked_object = image_pil * mask
-#     background = image_pil * ~mask
-
-
-#     '''Shifts image by x_units and y_units'''
-#     M = np.float32([[1, 0, x_units], [0, 1, y_units]])
-#     shifted_image = cv2.warpAffine(masked_object, M, (masked_object.shape[1] , masked_object.shape[0] ), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))
-#     masked_shifted_image = np.where(shifted_image[:, :, 0] != 0, True, False)
-
-#     '''Load stable diffuser model at checkpoint finetuned for inpainting task'''
-#     pipe = StableDiffusionInpaintPipeline.from_pretrained(
-#     #     "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16
-#         "stabilityai/stable-diffusion-2-inpainting",torch_dtype=torch.float16)
-#     pipe.to(CFG.device)
-
-
-#     # With Dilation
-#     structuring_element = np.ones((15, 15, 1), dtype=bool)
-#     extrapolated_mask = binary_dilation(mask, structure=structuring_element)
-#     mask_as_uint8 = extrapolated_mask.astype(np.uint8) * 255
-#     pil_mask = Image.fromarray(mask_as_uint8.squeeze(), mode='L').resize((1024, 1024))
-
-#     # # Without Dilation
-#     # pil_background = Image.fromarray(background)
-#     # mask_as_uint8 = mask.astype(np.uint8) * 255
-#     # pil_mask = Image.fromarray(mask_as_uint8.squeeze(), mode='L')
-
-#     '''Do inpainting on masked locations of original image'''
-#     prompt = 'fill as per background and neighborhood'
-#     inpainted_image = pipe(prompt=prompt, image=image_pil, mask_image=pil_mask).images[0]
-#     # inpainted_image
-
-#     '''Get composite of shifted object and background inpainted imaage'''
-#     pil_shifted_image = Image.fromarray(shifted_image).resize(inpainted_image.size)
-#     np_shifted_image = np.array(pil_shifted_image)
-#     masked_shifted_image = np.where(np_shifted_image[:, :, 0] != 0, True, False)
-#     masked_shifted_image = np.expand_dims(masked_shifted_image, axis=-1)
-#     inpainted_shifted = np.array(inpainted_image) * ~masked_shifted_image
-
-#     shifted_image = cv2.resize(shifted_image, inpainted_image.size)
-#     output = inpainted_shifted + shifted_image
-#     output = Image.fromarray(output)
-# #     visualize_results(image_pil, output, 'shifted')
+    '''Get masked object and background as two separate images'''
+    mask = np.expand_dims(masks[0], axis=-1)
+    masked_object = image_pil * mask
+    background = image_pil * ~mask
+
+
+    '''Shifts image by x_units and y_units'''
+    M = np.float32([[1, 0, x_units], [0, 1, y_units]])
+    shifted_image = cv2.warpAffine(masked_object, M, (masked_object.shape[1] , masked_object.shape[0] ), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))
+    masked_shifted_image = np.where(shifted_image[:, :, 0] != 0, True, False)
+
+    '''Load stable diffuser model at checkpoint finetuned for inpainting task'''
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    #     "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16
+        "stabilityai/stable-diffusion-2-inpainting",torch_dtype=torch.float16)
+    pipe.to(CFG.device)
+
+
+    # With Dilation
+    structuring_element = np.ones((15, 15, 1), dtype=bool)
+    extrapolated_mask = binary_dilation(mask, structure=structuring_element)
+    mask_as_uint8 = extrapolated_mask.astype(np.uint8) * 255
+    pil_mask = Image.fromarray(mask_as_uint8.squeeze(), mode='L').resize((1024, 1024))
+
+    # # Without Dilation
+    # pil_background = Image.fromarray(background)
+    # mask_as_uint8 = mask.astype(np.uint8) * 255
+    # pil_mask = Image.fromarray(mask_as_uint8.squeeze(), mode='L')
+
+    '''Do inpainting on masked locations of original image'''
+    prompt = 'fill as per background and neighborhood'
+    inpainted_image = pipe(prompt=prompt, image=image_pil, mask_image=pil_mask).images[0]
+    # inpainted_image
+
+    '''Get composite of shifted object and background inpainted imaage'''
+    pil_shifted_image = Image.fromarray(shifted_image).resize(inpainted_image.size)
+    np_shifted_image = np.array(pil_shifted_image)
+    masked_shifted_image = np.where(np_shifted_image[:, :, 0] != 0, True, False)
+    masked_shifted_image = np.expand_dims(masked_shifted_image, axis=-1)
+    inpainted_shifted = np.array(inpainted_image) * ~masked_shifted_image
+
+    shifted_image = cv2.resize(shifted_image, inpainted_image.size)
+    output = inpainted_shifted + shifted_image
+    output = Image.fromarray(output)
+#     visualize_results(image_pil, output, 'shifted')
     
     return output
 
 import gradio as gr
 
-with gr.Blocks() as demo:
-    gr.Markdown("Segmentation and shift")
-    with gr.Tab("Txt2Img"):
-        with gr.Row():
-            with gr.Column(scale=1):
-#                 image_pil, x_units, y_units, text_prompt
-
-                text_prompt = gr.Textbox(lines=1, label="Prompt")
-                image_pil = gr.Image(type="filepath", label='Image')
-#                 negative_prompt = gr.Textbox(lines=1, label="Negative Prompt")
-#                 width = gr.Dropdown(choices=number_choices, value=704, label="Width")
-#                 height = gr.Dropdown(choices=number_choicess, value=408, label="Height")
-                x_units = gr.Slider(minimum=50, maximum=300, step=1, value=10, label="x_units")
-                y_units = gr.Slider(minimum=30, maximum=300, step=0.1, value=5, label="y_units")
-#                 seed = gr.Textbox(label="Seed (Leave empty for random seed)")
-            with gr.Column(scale=2):
-                output = gr.Image(type="pil", label="Output Image")
-#                 generate = gr.Button("Generate")
-#                 output_seed = gr.Textbox(label="Current Seed")
-
-    # Create the txt2img function
+# with gr.Blocks() as demo:
+#     gr.Markdown("Segmentation and shift")
+#     with gr.Tab("Txt2Img"):
+#         with gr.Row():
+#             with gr.Column(scale=1):
+# #                 image_pil, x_units, y_units, text_prompt
+
+#                 text_prompt = gr.Textbox(lines=1, label="Prompt")
+#                 image_pil = gr.Image(type="filepath", label='Image')
+# #                 negative_prompt = gr.Textbox(lines=1, label="Negative Prompt")
+# #                 width = gr.Dropdown(choices=number_choices, value=704, label="Width")
+# #                 height = gr.Dropdown(choices=number_choicess, value=408, label="Height")
+#                 x_units = gr.Slider(minimum=50, maximum=300, step=1, value=10, label="x_units")
+#                 y_units = gr.Slider(minimum=30, maximum=300, step=0.1, value=5, label="y_units")
+# #                 seed = gr.Textbox(label="Seed (Leave empty for random seed)")
+#             with gr.Column(scale=2):
+#                 output = gr.Image(type="pil", label="Output Image")
+# #                 generate = gr.Button("Generate")
+# #                 output_seed = gr.Textbox(label="Current Seed")
+
+#     # Create the txt2img function
     
-#     generate.click(fn=main_fun, inputs=[image_pil, x_units, y_units, text_prompt], outputs=[output_image])
+# #     generate.click(fn=main_fun, inputs=[image_pil, x_units, y_units, text_prompt], outputs=[output_image])
+
+x_units = gr.Slider(minimum=50, maximum=300, step=1, value=10, label="x_units")
+y_units = gr.Slider(minimum=30, maximum=300, step=0.1, value=5, label="y_units")
+text_prompt = gr.Textbox(lines=1, label="Prompt")
+image_pil = gr.components.Image(type="filepath", label="Input Image")
+output = gr.components.Image(type="numpy", label="Output Image")
+]
 
 # Launch the Gradio UI
 gr.Interface(fn=main_fun, inputs=[image_pil, x_units, y_units, text_prompt], outputs=[output]).launch(share=True, debug=True)