87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
import cv2
|
|
import sys
|
|
import numpy as np
|
|
import onnxruntime as rt
|
|
import os
|
|
|
|
class Matting:
|
|
def __init__(self, model_path, input_size=(512, 512)):
|
|
# current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
# self.model_path = os.path.join(current_dir, "../checkpoint/modnet.onnx")
|
|
self.sess = rt.InferenceSession(model_path)
|
|
self.input_name = self.sess.get_inputs()[0].name
|
|
self.label_name = self.sess.get_outputs()[0].name
|
|
self.input_size = input_size
|
|
self.txt_font = cv2.FONT_HERSHEY_PLAIN
|
|
|
|
def normalize(self, im, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
|
|
im = im.astype(np.float32, copy=False) / 255.0
|
|
im -= mean
|
|
im /= std
|
|
return im
|
|
|
|
def resize(self, im, target_size=608, interp=cv2.INTER_LINEAR):
|
|
if isinstance(target_size, list) or isinstance(target_size, tuple):
|
|
w = target_size[0]
|
|
h = target_size[1]
|
|
else:
|
|
w = target_size
|
|
h = target_size
|
|
im = cv2.resize(im, (w, h), interpolation=interp)
|
|
return im
|
|
|
|
def preprocess(self, image, target_size=(512, 512), interp=cv2.INTER_LINEAR):
|
|
image = self.normalize(image)
|
|
image = self.resize(image, target_size=target_size, interp=interp)
|
|
image = np.transpose(image, [2, 0, 1])
|
|
image = image[None, :, :, :]
|
|
return image
|
|
|
|
def predict_frame(self, bgr_image):
|
|
assert len(bgr_image.shape) == 3, "Please input RGB image."
|
|
raw_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
|
|
h, w, c = raw_image.shape
|
|
image = self.preprocess(raw_image, target_size=self.input_size)
|
|
|
|
pred = self.sess.run(
|
|
[self.label_name],
|
|
{self.input_name: image.astype(np.float32)}
|
|
)[0]
|
|
pred = pred[0, 0]
|
|
matte_np = self.resize(pred, target_size=(w, h), interp=cv2.INTER_NEAREST)
|
|
matte_np = np.expand_dims(matte_np, axis=-1)
|
|
return matte_np
|
|
|
|
def predict_image(self, source_image_path, save_image_path):
|
|
bgr_image = cv2.imread(source_image_path)
|
|
assert len(bgr_image.shape) == 3, "Please input RGB image."
|
|
matte_np = self.predict_frame(bgr_image)
|
|
# bgr_image = np.ones_like(bgr_image)*255
|
|
matting_frame = matte_np * bgr_image + (1 - matte_np) * np.full(bgr_image.shape, 0.0)
|
|
matting_frame = matting_frame.astype('uint8')
|
|
cv2.imwrite(save_image_path, matting_frame)
|
|
def predict_image_rgba(self, source_image_path, save_image_path):
|
|
bgr_image = cv2.imread(source_image_path)
|
|
assert len(bgr_image.shape) == 3, "Please input RGB image."
|
|
matte_np = self.predict_frame(bgr_image)
|
|
# bgr_image = np.ones_like(bgr_image)*255
|
|
matting_frame = matte_np * bgr_image + (1 - matte_np) * np.full(bgr_image.shape, 0.0)
|
|
matte_np = (matte_np*255).astype('uint8')
|
|
matting_frame = matting_frame.astype('uint8')
|
|
matting_frame = np.concatenate([matting_frame, matte_np],axis = 2)
|
|
cv2.imwrite(save_image_path, matting_frame)
|
|
|
|
def main():
|
|
# 检查命令行参数的数量
|
|
if len(sys.argv) != 3:
|
|
print("Usage: python interface/matting.py <img_path> <output_path>")
|
|
sys.exit(1) # 参数数量不正确时退出程序
|
|
|
|
img_path = sys.argv[1]
|
|
output_path = sys.argv[2]
|
|
print(f"img path is set to: {img_path}, output path is set to: {output_path}")
|
|
model = Matting()
|
|
model.predict_image_rgba(img_path, output_path)
|
|
|
|
if __name__ == "__main__":
|
|
main() |