Update app.py

app.py

@@ -11,8 +11,8 @@ os.mkdir("data")
 os.rename("best.ckpt", "models/best.ckpt")
 os.mkdir("dataout")
 model = hub.Module(name='U2Net')
-def infer(img,mask,option):
-  img = ImageOps.contain(img, (700,700))
+def infer(img,option):
+  img = ImageOps.contain(img["image"], (700,700))
   width, height = img.size
   img.save("./data/data.png")
   if option == "automatic (U2net)":
@@ -25,13 +25,13 @@ def infer(img,mask,option):
           visualization=True)
       im = Image.fromarray(result[0]['mask'])
   else:
-      mask = mask.resize((width,height))
+      mask = img["mask"].resize((width,height))
       im = mask
   im.save("./data/data_mask.png")
   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",im
   
-inputs = [gr.inputs.Image(type='pil', label="Original Image"),gr.inputs.Image(type='pil',source="canvas", label="Mask",invert_colors=True),gr.inputs.Radio(choices=["automatic (U2net)","manual"], type="value", default="manual", label="Masking option")]
+inputs = [gr.inputs.Image(type='pil',tool="sketch", label="Input"),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="pil",label="Mask")]
 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"