Add application file

Weight/sam_vit_b_01ec64.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912
+size 375042383

Weight/sam_vit_h_4b8939.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e
+size 2564550879

Weight/sam_vit_l_0b3195.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622
+size 1249524607

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+import numpy as np
+import torch
+from diffusers import StableDiffusionInpaintPipeline
+from PIL import Image
+from segment_anything import SamPredictor, sam_model_registry
+
+
+device="cpu"
+sam_checkpoint = "Weight/sam_vit_h_4b8939.pth"
+model_type = "vit_h"
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+sam.to(device)
+predictor = SamPredictor(sam)
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-inpainting",
+    torch_dtype=torch.float32
+)
+
+pipe = pipe.to(device)
+selected_pixels = []
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        input_img = gr.Image(label="Input")
+        mask_img = gr.Image(label="Mas")
+        output_img = gr.Image(label="Output")
+    with gr.Blocks():
+        prompt_text = gr.Textbox(lines=1, label="Prompt")
+    with gr.Blocks():
+        submit = gr.Button("Submit")
+
+    def generate_mask(image, evt:gr.SelectData):
+      
+      input_labels = np.ones(len(selected_pixels))
+      selected_pixels.append(evt.index)
+
+      predictor.set_image(image)
+      input_points = np.array(selected_pixels)
+
+      input_labels = np.ones(input_labels.shape[0])
+  
+      mask, _, _ = predictor.predict(
+          point_coords= input_points,
+          point_labels= input_labels,
+          multimask_output=False
+      )
+      # (n, sz, sz)
+      mask = Image.fromarray(mask[0, : , :])
+      mask = mask.resize((512, 512)) # Resize the mask to (512, 512)
+      mask = np.expand_dims(mask, axis=2)
+      return mask
+
+    def inpaint(image, mask, prompt):
+        image = Image.fromarray(image)
+        mask = Image.fromarray(mask)
+
+        image = image.resize((512,512))
+        mask = mask.resize((512,512))
+
+        output = pipe(
+            prompt=prompt, 
+            image=image, 
+            mask_image=mask,
+            ).images[0]
+            
+        return output
+    
+    input_img.select(generate_mask, [input_img], [mask_img])
+    submit.click(
+        inpaint,
+        inputs=[input_img, mask_img, prompt_text],
+        outputs=[output_img],
+    )
+if __name__ == "__main__":
+    demo.launch()