Update app.py

app.py

@@ -1,90 +1,139 @@
 import os
-import imageio
-from PIL import Image
+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')
 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)
 
-# 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)
+# 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!")
 
-# Load LaMa ONNX model
-sess_options = onnxruntime.SessionOptions()
-lama_model = onnxruntime.InferenceSession('lama_fp32.onnx', sess_options=sess_options)
+    if img.ndim == 3:
+        img = np.transpose(img, (2, 0, 1))  # chw
+    elif img.ndim == 2:
+        img = img[np.newaxis, ...]
 
-# Load U^2-Net model for automatic masking
-u2net_model = hub.Module(name='U2Net') 
+    assert img.ndim == 3
 
-# --- Helper Functions ---
+    img = img.astype(np.float32) / 255
+    return img
 
-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)
+
+def ceil_modulo(x, mod):
+    if x % mod == 0:
+        return x
+    return (x // mod + 1) * mod
+
+
+def scale_image(img, factor, interpolation=cv2.INTER_AREA):
+    if img.shape[0] == 1:
+        img = img[0]
     else:
-        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})
+        img = np.transpose(img, (1, 2, 0))
+
+    img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
+
+    if img.ndim == 2:
+        img = img[None, ...]
+    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)})
+
     output = outputs[0][0]
+    # Postprocess the outputs
     output = output.transpose(1, 2, 0)
-    output = (output * 255).astype(np.uint8)
-    return Image.fromarray(output) 
+    output = output.astype(np.uint8)
+    output = Image.fromarray(output)
+    output = output.resize(imagex.size)
+    output.save("/home/user/app/dataout/data_mask.png")
+
 
-# --- Gradio Interface ---
+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"
 
-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,
+    fn=infer,
     inputs=[
         gr.Image(label="Input Image", type="numpy"),
         gr.Radio(choices=["automatic", ],