Update app.py

app.py

@@ -1,40 +1,72 @@
 import os
-os.system("wget https://huggingface.co/akhaliq/lama/resolve/main/best.ckpt")
 os.system("pip install imageio")
-import cv2
-import paddlehub as hub
+# app.py for Hugging Face Spaces
+
 import gradio as gr
-import torch
-from PIL import Image, ImageOps
+from simple_lama_inpainting import SimpleLama
+from PIL import Image
 import numpy as np
-import imageio
-os.mkdir("data")
-os.rename("best.ckpt", "models/best.ckpt")
-os.mkdir("dataout")
-model = hub.Module(name='U2Net')
-def infer(img,option):
-  print(type(img))
-  print(type(img["image"]))
-  print(type(img["mask"]))
-  imageio.imwrite("./data/data.png", img["image"])
-  if option == "automatic (U2net)":
-      result = model.Segmentation(
-          images=[cv2.cvtColor(img["image"], cv2.COLOR_RGB2BGR)],
-          paths=None,
-          batch_size=1,
-          input_size=320,
-          output_dir='output',
-          visualization=True)
-      im = Image.fromarray(result[0]['mask'])
-      im.save("./data/data_mask.png")
-  else:
-      imageio.imwrite("./data/data_mask.png", img["mask"])
-  os.system('python predict.py model.path=/home/user/app/ indir=/home/user/app/data/ outdir=/home/user/app/dataout/ device=cpu')
-  return "./dataout/data_mask.png","./data/data_mask.png"
-  
-inputs = [gr.Image(tool="sketch", label="Input",type="numpy"),gr.inputs.Radio(choices=["automatic (U2net)","manual"], type="value", default="manual", label="Masking option")]
-outputs = [gr.outputs.Image(type="file",label="output"),gr.outputs.Image(type="file",label="Mask")]
+import os
+
+# Initialize the SimpleLama model
+simple_lama = SimpleLama()
+
+def inpaint_image(input_data):
+    if input_data is None:
+        raise gr.Error("Please upload an image and draw a mask.")
+    
+    # input_data is a dict with 'image' and 'mask' as paths or arrays
+    # In Gradio 4+, it's dict with 'image' and 'mask' as PIL Images or paths
+    image_path = input_data['image']
+    mask_path = input_data['mask']
+    
+    # Open image and mask
+    image = Image.open(image_path).convert("RGB")
+    mask = Image.open(mask_path).convert("L")
+    
+    # In Gradio sketch, mask is white where drawn (to inpaint), black elsewhere
+    # But ensure it's binary if needed, though SimpleLama likely handles grayscale
+    
+    # Perform inpainting
+    result = simple_lama(image, mask)
+    
+    return result
+
+with gr.Blocks(title="Image Inpainting with SimpleLama") as demo:
+    gr.Markdown("# Image Inpainting App")
+    gr.Markdown("Upload an image, draw a mask on the areas to inpaint, and click Submit.")
+    
+    with gr.Row():
+        input_image = gr.Image(
+            label="Upload and Mask Image",
+            tool="sketch",
+            type="filepath",  # Returns paths
+            brush_color="#FFFFFF",  # White brush for mask
+            interactive=True
+        )
+        output_image = gr.Image(label="Inpainted Image")
+    
+    submit_btn = gr.Button("Submit")
+    submit_btn.click(inpaint_image, inputs=[input_image], outputs=[output_image])
+    
+    gr.Examples(
+        examples=[
+            # ["examples/example1.jpg"],
+            ["person512.png"],
+        ],
+        inputs=[input_image],
+        label="Try these examples (upload and draw mask)"
+    )
+
+# For HF Spaces, requirements.txt would include:
+# simple-lama-inpainting
+# gradio
+# pillow
+# torch (if not included in simple-lama)
+
+if __name__ == "__main__":
+    demo.launch()
 title = "LaMa Image Inpainting"
 description = "Gradio demo for LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Masks are generated by U^2net"
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.07161' target='_blank'>Resolution-robust Large Mask Inpainting with Fourier Convolutions</a> | <a href='https://github.com/saic-mdal/lama' target='_blank'>Github Repo</a></p>"
-gr.Interface(infer, inputs, outputs, title=title, description=description, article=article).launch()
+# gr.Interface(infer, inputs, outputs, title=title, description=description, article=article).launch()