Update app.py

app.py

@@ -20,9 +20,11 @@ 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, drawing):
-    """ Process the input image and drawing using SAM model and processor. """
-    inputs = processor(image, return_tensors="pt").to(device)
+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)
     with torch.no_grad():
         outputs = model(**inputs)
     masks = processor.image_processor.post_process_masks(
@@ -53,10 +55,10 @@ def inpaint(raw_image, input_mask, prompt, negative_prompt=None, seed=74294536,
     ).images[0]
     return image
 
-def gradio_interface(image, drawing, positive_prompt, negative_prompt):
-    """ Gradio interface function to handle image, drawing, and prompts. """
+def gradio_interface(image, annotation, positive_prompt, negative_prompt):
+    """ Gradio interface function to handle image, annotated drawing, and prompts. """
     raw_image = Image.fromarray(image).convert("RGB").resize((512, 512))
-    mask = get_processed_inputs(raw_image, drawing)
+    mask = get_processed_inputs(raw_image, annotation)
     processed_image = inpaint(raw_image, mask, positive_prompt, negative_prompt)
     return processed_image, mask_to_rgb(mask)
 
@@ -64,7 +66,7 @@ iface = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Image(type="numpy", label="Input Image"),
-        gr.Sketch(label="Draw on the image", shape=(512, 512)),
+        gr.Image(tool="editor", label="Draw on the image", output="png", shape=(512, 512)),
         gr.Textbox(label="Positive Prompt", placeholder="Enter positive prompt here"),
         gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here")
     ],