chore: inference python scripts

.gitignore

@@ -37,3 +37,7 @@ runs/
 # model checkpoints
 models/
 .env
+
+# experiments
+BallonsTranslator/
+carve-lama/

scripts/ctd_inference.py

@@ -0,0 +1,31 @@
+import numpy as np
+import cv2
+
+img = cv2.imread('data/1746025823_segment.png')
+
+kernel = np.ones((3,3),np.uint8)
+h, w = img.shape[0], img.shape[1]
+seedpnt = (int(w/2), int(h/2))
+difres = 10
+
+# convert to grayscale
+img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+# ballon_mask = img - 127
+ballon_mask = 127 - img
+ballon_mask = img
+ballon_mask = cv2.dilate(ballon_mask, kernel,iterations = 1)
+# ballon_area, _, _, rect = cv2.floodFill(ballon_mask, mask=None, seedPoint=seedpnt,  flags=4, newVal=(30), loDiff=(difres, difres, difres), upDiff=(difres, difres, difres))
+ballon_mask = 30 - ballon_mask
+retval, ballon_mask = cv2.threshold(ballon_mask, 1, 255, cv2.THRESH_BINARY)
+ballon_mask = cv2.bitwise_not(ballon_mask, ballon_mask)
+
+# box_kernel = int(np.sqrt(ballon_area) / 30)
+# if box_kernel > 1:
+#     box_kernel = np.ones((box_kernel,box_kernel),np.uint8)
+#     ballon_mask = cv2.dilate(ballon_mask, box_kernel, iterations = 1)
+#     ballon_mask = cv2.erode(ballon_mask, box_kernel, iterations = 1)
+
+cv2.imshow('ballon_mask', ballon_mask)
+#cv2.imshow('img', img)
+cv2.waitKey(0)

scripts/inference_inpaint_onnx.py

@@ -0,0 +1,111 @@
+import cv2
+import numpy as np
+import onnxruntime
+import torch
+import io
+import requests
+from PIL import Image
+
+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
+
+
+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:
+        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 open_image(image):
+    if isinstance(image, str):
+      if image.startswith("http://") or image.startswith("https://"):
+        image = Image.open(io.BytesIO(requests.get(image).content))
+      else:
+        image = Image.open(image)
+    return image
+
+
+
+sess_options = onnxruntime.SessionOptions()
+model = onnxruntime.InferenceSession('models/lama_manga.onnx', sess_options=sess_options)
+
+image_url = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/image.jpg" # @param {type:"string"}
+mask_url = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/mask.png" # @param {type:"string"}
+
+image = open_image(image_url).resize((512, 512))
+mask = open_image(mask_url).convert("L").resize((512, 512))
+
+image, mask = prepare_img_and_mask(image, mask, 'cpu')
+# Run the model
+outputs = model.run(None,
+                    {'image': image.numpy().astype(np.float32),
+                     'mask': mask.numpy().astype(np.float32)})
+
+
+output = outputs[0][0] * 256
+# Postprocess the outputs
+output = output.transpose(1, 2, 0)
+output = output.astype(np.uint8)
+output = Image.fromarray(output)
+output.show()