Update app.py

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 import numpy as np
 from transformers import SamModel, SamProcessor
 from diffusers import AutoPipelineForInpainting
+from diffusers.models.autoencoders.vq_model import VQEncoderOutput, VQModel
 import torch
 
 # Check if GPU is available, otherwise use CPU
@@ -15,16 +16,14 @@ model = SamModel.from_pretrained(model_name).to(device)
 processor = SamProcessor.from_pretrained(model_name)
 
 def mask_to_rgb(mask):
-    """ Convert binary mask to RGB with transparency for the background. """
     bg_transparent = np.zeros(mask.shape + (4,), dtype=np.uint8)
-    bg_transparent[mask == 1] = [0, 255, 0, 127]  # Green mask with some transparency
+    bg_transparent[mask == 1] = [0, 255, 0, 127]
     return bg_transparent
 
-def get_processed_inputs(image, annotation):
-    """ Process the input image and annotated drawing using SAM model and processor. """
-    mask = np.zeros(image.size, dtype=np.uint8)
-    mask[annotation[:,:,3] > 128] = 1  # Assume drawing is in alpha channel of RGBA
-    inputs = processor(images=image, return_tensors="pt").to(device)
+def get_processed_inputs(image, points_str):
+    points = [list(map(int, point.split(','))) for point in points_str.split()]
+    input_points = [points]
+    inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
     masks = processor.image_processor.post_process_masks(
@@ -36,7 +35,6 @@ def get_processed_inputs(image, annotation):
     return ~best_mask.cpu().numpy()
 
 def inpaint(raw_image, input_mask, prompt, negative_prompt=None, seed=74294536, cfgs=7):
-    """ Inpaint the masked area in the image using a text prompt and an inpainting pipeline. """
     mask_image = Image.fromarray(input_mask)
     rand_gen = torch.manual_seed(seed)
     pipeline = AutoPipelineForInpainting.from_pretrained(
@@ -55,10 +53,9 @@ def inpaint(raw_image, input_mask, prompt, negative_prompt=None, seed=74294536,
     ).images[0]
     return image
 
-def gradio_interface(image, annotation, positive_prompt, negative_prompt):
-    """ Gradio interface function to handle image, annotated drawing, and prompts. """
+def gradio_interface(image, points, positive_prompt, negative_prompt):
     raw_image = Image.fromarray(image).convert("RGB").resize((512, 512))
-    mask = get_processed_inputs(raw_image, annotation)
+    mask = get_processed_inputs(raw_image, points)
     processed_image = inpaint(raw_image, mask, positive_prompt, negative_prompt)
     return processed_image, mask_to_rgb(mask)
 
@@ -66,7 +63,7 @@ iface = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Image(type="numpy", label="Input Image"),
-        gr.Image(tool="editor", label="Draw on the image", output="png", shape=(512, 512)),
+        gr.Textbox(label="Points (format: x1,y1 x2,y2 ...)", placeholder="e.g., 100,100 200,200"),
         gr.Textbox(label="Positive Prompt", placeholder="Enter positive prompt here"),
         gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here")
     ],
@@ -75,7 +72,7 @@ iface = gr.Interface(
         gr.Image(label="Segmentation Mask")
     ],
     title="Interactive Image Inpainting",
-    description="Draw on the image to select areas for segmentation, provide prompts, and see the inpainted result."
+    description="Enter points as 'x1,y1 x2,y2 ...' for segmentation, provide prompts, and see the inpainted result."
 )
 
 iface.launch(share=True)