Update app.py

app.py

@@ -1,138 +1,100 @@
 import os
-os.system("wget https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx")
-os.system("pip install onnxruntime imageio")
-import cv2
-import paddlehub as hub
-import gradio as gr
-import torch
-from PIL import Image, ImageOps
-import numpy as np
 import imageio
-os.mkdir("data")
-os.mkdir("dataout")
-model = hub.Module(name='U2Net')
+from PIL import Image
+import gradio as gr
 import cv2
-import numpy as np
+import paddlehub as hub
 import onnxruntime
-import torch
-from PIL import Image
-sess_options = onnxruntime.SessionOptions()
-rmodel = onnxruntime.InferenceSession('lama_fp32.onnx', sess_options=sess_options)
-
-# Source https://github.com/advimman/lama
-def get_image(image):
-    if isinstance(image, Image.Image):
-        img = np.array(image)
-    elif isinstance(image, np.ndarray):
-        img = image.copy()
-    else:
-        raise Exception("Input image should be either PIL Image or numpy array!")
-
-    if img.ndim == 3:
-        img = np.transpose(img, (2, 0, 1))  # chw
-    elif img.ndim == 2:
-        img = img[np.newaxis, ...]
-
-    assert img.ndim == 3
-
-    img = img.astype(np.float32) / 255
-    return img
 
+# Download and setup models
+os.system("wget https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx")
+os.system("pip install onnxruntime imageio")
+os.makedirs("data", exist_ok=True)
+os.makedirs("dataout", exist_ok=True)
 
-def ceil_modulo(x, mod):
-    if x % mod == 0:
-        return x
-    return (x // mod + 1) * mod
+# Load LaMa ONNX model
+sess_options = onnxruntime.SessionOptions()
+lama_model = onnxruntime.InferenceSession('lama_fp32.onnx', sess_options=sess_options)
 
+# Load U^2-Net model for automatic masking
+u2net_model = hub.Module(name='U2Net') 
 
-def scale_image(img, factor, interpolation=cv2.INTER_AREA):
-    if img.shape[0] == 1:
-        img = img[0]
-    else:
-        img = np.transpose(img, (1, 2, 0))
+# --- Helper Functions ---
 
-    img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
-
-    if img.ndim == 2:
-        img = img[None, ...]
+def prepare_image(image, target_size=(512, 512)):
+    """Resizes and preprocesses image for LaMa model."""
+    if isinstance(image, Image.Image):
+        image = image.resize(target_size)
+        image = np.array(image)
+    elif isinstance(image, np.ndarray):
+        image = cv2.resize(image, target_size)
     else:
-        img = np.transpose(img, (2, 0, 1))
-    return img
-
-
-def pad_img_to_modulo(img, mod):
-    channels, height, width = img.shape
-    out_height = ceil_modulo(height, mod)
-    out_width = ceil_modulo(width, mod)
-    return np.pad(
-        img,
-        ((0, 0), (0, out_height - height), (0, out_width - width)),
-        mode="symmetric",
-    )
-
-
-def prepare_img_and_mask(image, mask, device, pad_out_to_modulo=8, scale_factor=None):
-    out_image = get_image(image)
-    out_mask = get_image(mask)
-
-    if scale_factor is not None:
-        out_image = scale_image(out_image, scale_factor)
-        out_mask = scale_image(out_mask, scale_factor, interpolation=cv2.INTER_NEAREST)
-
-    if pad_out_to_modulo is not None and pad_out_to_modulo > 1:
-        out_image = pad_img_to_modulo(out_image, pad_out_to_modulo)
-        out_mask = pad_img_to_modulo(out_mask, pad_out_to_modulo)
-
-    out_image = torch.from_numpy(out_image).unsqueeze(0).to(device)
-    out_mask = torch.from_numpy(out_mask).unsqueeze(0).to(device)
-
-    out_mask = (out_mask > 0) * 1
-
-    return out_image, out_mask
-
-
-def predict(jpg, msk):
-
-
-    imagex = Image.open(jpg)
-    mask = Image.open(msk).convert("L")
-
-    image, mask = prepare_img_and_mask(imagex.resize((512, 512)), mask.resize((512, 512)), 'cpu')
-    # Run the model
-    outputs = rmodel.run(None, {'image': image.numpy().astype(np.float32), 'mask': mask.numpy().astype(np.float32)})
-
+        raise ValueError("Input image should be either PIL Image or numpy array!")
+
+    # Normalize to [0, 1] and convert to CHW format
+    image = image.astype(np.float32) / 255.0
+    if image.ndim == 3:
+        image = np.transpose(image, (2, 0, 1))
+    elif image.ndim == 2:
+        image = image[np.newaxis, ...]
+    return image[np.newaxis, ...] # Add batch dimension
+
+def generate_mask(image, method="automatic"):
+    """Generates mask from image using U^2-Net or user input."""
+    if method == "automatic":
+        input_size = 320  # Adjust based on U^2-Net requirements
+        result = u2net_model.Segmentation(
+            images=[cv2.cvtColor(image, cv2.COLOR_RGB2BGR)],
+            paths=None,
+            batch_size=1,
+            input_size=input_size,
+            output_dir='output',
+            visualization=False
+        )
+        mask = Image.fromarray(result[0]['mask'])
+        mask = mask.resize((512, 512))  # Resize to match LaMa input
+        mask.save("./data/data_mask.png")
+    else:  # "manual"
+        mask = imageio.imread("./data/data_mask.png") 
+        mask = Image.fromarray(mask).convert("L")  # Ensure grayscale
+        mask = mask.resize((512, 512))
+    return prepare_image(mask, (512, 512))
+
+def inpaint_image(image, mask):
+    """Performs inpainting using the LaMa model."""
+    outputs = lama_model.run(None, {'image': image, 'mask': mask})
     output = outputs[0][0]
-    # Postprocess the outputs
     output = output.transpose(1, 2, 0)
-    output = output.astype(np.uint8)
-    output = Image.fromarray(output)
-    output = output.resize(imagex.size)
-    output.save("/home/user/app/dataout/data_mask.png")
-
-
-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"])
-  predict("./data/data.png", "./data/data_mask.png")    
-  return "./dataout/data_mask.png","./data/data_mask.png"
-  
-inputs = [gr.Image(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")]
-title = "LaMa Image Inpainting (using ONNX model from Carve))"
-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://huggingface.co/Carve/LaMa-ONNX' target='_blank'>ONNX model ported by Carve.Photos</a> | <a href='https://github.com/saic-mdal/lama' target='_blank'>LaMa github repo</a></p>"
-gr.Interface(infer, inputs, outputs, title=title, description=description, article=article).launch()
+    output = (output * 255).astype(np.uint8)
+    return Image.fromarray(output) 
+
+# --- Gradio Interface ---
+
+def process_image(input_image, mask_option):
+    """Main function for Gradio interface."""
+    imageio.imwrite("./data/data.png", input_image)
+
+    image = prepare_image(input_image)
+    mask = generate_mask(input_image, method=mask_option)
+    
+    inpainted_image = inpaint_image(image, mask)
+    inpainted_image = inpainted_image.resize(Image.open("./data/data.png").size)
+    inpainted_image.save("./dataout/data_mask.png")
+    return "./dataout/data_mask.png", "./data/data_mask.png"
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(label="Input Image", type="numpy"),
+        gr.Radio(choices=["automatic", "manual"], 
+                 type="value", default="manual", label="Masking Option")
+    ],
+    outputs=[
+        gr.Image(type="file", label="Inpainted Image"),
+        gr.Image(type="file", label="Generated Mask")
+    ],
+    title="LaMa Image Inpainting",
+    description="Image inpainting with LaMa and U^2-Net. Upload your image and choose automatic or manual masking.",
+)
+
+iface.launch()