Update app.py

app.py

@@ -20,9 +20,9 @@ def mask_to_rgb(mask):
     bg_transparent[mask == 1] = [0, 255, 0, 127]  # Green mask with some transparency
     return bg_transparent
 
-def get_processed_inputs(image, points):
-    """ Process the input image and points using SAM model and processor. """
-    inputs = processor(image, input_points=points, return_tensors="pt").to(device)
+def get_processed_inputs(image, drawing):
+    """ Process the input image and drawing using SAM model and processor. """
+    inputs = processor(image, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
     masks = processor.image_processor.post_process_masks(
@@ -53,11 +53,10 @@ def inpaint(raw_image, input_mask, prompt, negative_prompt=None, seed=74294536,
     ).images[0]
     return image
 
-def gradio_interface(image, points_json, positive_prompt, negative_prompt):
-    """ Gradio interface function to handle image, points for segmentation, and prompts. """
-    points = [[(point['x'], point['y']) for point in stroke['points']] for stroke in points_json]
+def gradio_interface(image, drawing, positive_prompt, negative_prompt):
+    """ Gradio interface function to handle image, drawing, and prompts. """
     raw_image = Image.fromarray(image).convert("RGB").resize((512, 512))
-    mask = get_processed_inputs(raw_image, points)
+    mask = get_processed_inputs(raw_image, drawing)
     processed_image = inpaint(raw_image, mask, positive_prompt, negative_prompt)
     return processed_image, mask_to_rgb(mask)
 
@@ -65,7 +64,7 @@ iface = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Image(type="numpy", label="Input Image"),
-        gr.Image(type="json", label="Click to select points", tool="sketch", brush_radius=1, shape=(512, 512)),
+        gr.Sketch(label="Draw on the image", shape=(512, 512)),
         gr.Textbox(label="Positive Prompt", placeholder="Enter positive prompt here"),
         gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here")
     ],
@@ -74,7 +73,7 @@ iface = gr.Interface(
         gr.Image(label="Segmentation Mask")
     ],
     title="Interactive Image Inpainting",
-    description="Click on the image to select points for segmentation, provide prompts, and see the inpainted result."
+    description="Draw on the image to select areas for segmentation, provide prompts, and see the inpainted result."
 )
 
-iface.launch(share=True)
+iface.launch(share=True