Update app.py

app.py

@@ -1,16 +1,22 @@
-import gradio as gr
-
 import os
-home = os.getcwd()
-home
+import warnings
+warnings.filterwarnings('ignore')
+
+import subprocess, io, os, sys, time
+
+# os.system("pip install -q gradio==3.48.0")
+os.system("pip install -q gradio")
+os.system("pip install -q diffusers")
+
+result = subprocess.run(['pip', 'install', '-e', 'GroundingDINO'], check=True)
+print(f'pip install GroundingDINO = {result}')
+sys.path.insert(0, './GroundingDINO')
+
+# text_prompt = 'basket'
+# image_path = '/kaggle/input/avataar/wall hanging.jpg'
+# output_image_path = '/kaggle/working'
 
-!git clone https://github.com/IDEA-Research/GroundingDINO
-%cd /{home}/GroundingDINO   
-!pip install -q -e .
 
-text_prompt = 'basket'
-image_path = '/kaggle/input/avataar/wall hanging.jpg'
-output_image_path = '/kaggle/working'
 
 
 '''Importing Libraries'''
@@ -27,6 +33,9 @@ from groundingdino.util.utils import clean_state_dict
 from huggingface_hub import hf_hub_download
 from segment_anything import sam_model_registry
 from segment_anything import SamPredictor
+from diffusers import StableDiffusionInpaintPipeline
+
+from scipy.ndimage import binary_dilation
 
 import cv2
 import matplotlib.pyplot as plt
@@ -101,8 +110,7 @@ def build_groundingdino():
     groundingdino = load_model_hf(ckpt_repo_id, ckpt_filename, ckpt_config_filename)
     return groundingdino
 
-model_sam = build_sam()
-model_groundingdino = build_groundingdino()
+
 
 
 '''Predictions'''
@@ -154,10 +162,6 @@ def draw_image(image_pil, masks, boxes, alpha=0.4):
         image = draw_segmentation_masks(image, masks=masks, colors=['red'] * len(masks), alpha=alpha)
     return image.numpy().transpose(1, 2, 0)
 
-image_pil = load_image(image_path)
-
-masks, boxes, phrases, logits = mask_predict(image_pil, text_prompt=text_prompt, box_threshold=0.23, text_threshold=0.25)
-output = draw_image(image_pil, masks, boxes, alpha=0.4)
 # torch.save(masks, 'masks.pt')
 
 
@@ -175,18 +179,25 @@ def visualize_results(img1, img2, task):
     for ax in axes:
         ax.axis('off')
 
-visualize_results(image_pil, output, 'segmented')
+# visualize_results(image_pil, output, 'segmented')
 
+# x_units = 200
+# y_units = -100
+# text_prompt = 'wooden stool'
+# image_path = '/kaggle/input/avataar/stool.jpeg'
+# output_image_path = '/kaggle/working'
 
-x_units = 200
-y_units = -100
-# import torch
-# import numpy as np
-# masks = torch.load('/kaggle/input/chair-mask/masks.pt')
-# print(masks.shape)
-# masks
+def main_fun(image_pil, x_units, y_units, text_prompt):
+#     x_units = 200
+#     y_units = -100
+#     text_prompt = 'wooden stool'
+    model_sam = build_sam()
+    model_groundingdino = build_groundingdino()
+    
+#     image_pil = load_image(image_path)
+    masks, boxes, phrases, logits = mask_predict(image_pil, text_prompt=text_prompt, box_threshold=0.23, text_threshold=0.25)
+    output = draw_image(image_pil, masks, boxes, alpha=0.4)
 
-def main_fun():
 
     '''Get masked object and background as two separate images'''
     mask = np.expand_dims(masks[0], axis=-1)
@@ -200,7 +211,6 @@ def main_fun():
     masked_shifted_image = np.where(shifted_image[:, :, 0] != 0, True, False)
 
     '''Load stable diffuser model at checkpoint finetuned for inpainting task'''
-    from diffusers import StableDiffusionInpaintPipeline
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
     #     "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16
         "stabilityai/stable-diffusion-2-inpainting",torch_dtype=torch.float16)
@@ -208,22 +218,18 @@ def main_fun():
 
 
     # With Dilation
-
-    from scipy.ndimage import binary_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))
-    # pil_mask
 
     # # 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')
-    # # pil_mask
 
     '''Do inpainting on masked locations of original image'''
-    prompt = 'a photo of background'
+    prompt = 'fill as per background and neighborhood'
     inpainted_image = pipe(prompt=prompt, image=image_pil, mask_image=pil_mask).images[0]
     # inpainted_image
 
@@ -237,21 +243,36 @@ def main_fun():
     shifted_image = cv2.resize(shifted_image, inpainted_image.size)
     output = inpainted_shifted + shifted_image
     output = Image.fromarray(output)
-    visualize_results(image_pil, output, 'shifted')
-
-
-
-inputs_image = [
-    gr.components.Image(type="filepath", label="Input Image"),
-]
-outputs_image = [
-    gr.components.Image(type="numpy", label="Output Image"),
-]
-interface_image = gr.Interface(
-    fn=main_fun,
-    inputs=inputs_image,
-    outputs=outputs_image,
-    title="Pothole detector",
-    # examples=path,
-    cache_examples=False,
-)
+#     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(label='Image Input', type='pil', tool='sketch', source="upload")
+#                 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_image = gr.Image(container=True, height=500, width=500)
+#                 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])
+
+# Launch the Gradio UI
+gr.Interface(fn=main_fun, inputs=[image_pil, x_units, y_units, text_prompt], outputs=[output_image]).launch(share=True, debug=True)
+# demo