main
John Lauer Publish 0.7.0 8ff960b 9h ago
"""Cut a person out of a photo: python3 cutout.py <photo.jpg> <out.png> [model.onnx]
Needs onnxruntime, numpy and Pillow, and the u2net_human_seg model (about 170 MB):
  curl -L -o u2net_human_seg.onnx https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx
Runs offline on the CPU in a few seconds."""
import sys, numpy as np, onnxruntime as ort
from PIL import Image, ImageFilter
def cut(src, dst, model='u2net_human_seg.onnx'):
    im=Image.open(src).convert('RGB'); W,H=im.size
    x=np.asarray(im.resize((320,320),Image.LANCZOS),dtype=np.float32); x/=x.max()
    x=(x-np.array([0.485,0.456,0.406],dtype=np.float32))/np.array([0.229,0.224,0.225],dtype=np.float32)
    x=x.transpose(2,0,1)[None].astype(np.float32)
    s=ort.InferenceSession(model,providers=['CPUExecutionProvider']); o=s.run(None,{s.get_inputs()[0].name:x})[0][0,0]
    o=(o-o.min())/(o.max()-o.min()+1e-8)
    m=Image.fromarray((o*255).astype(np.uint8)).resize((W,H),Image.BICUBIC).filter(ImageFilter.GaussianBlur(2))
    a=np.asarray(m,dtype=np.float32)/255; a=np.clip((a-0.35)/0.4,0,1)            # firm up the soft edge
    out=im.convert('RGBA'); out.putalpha(Image.fromarray((a*255).astype(np.uint8))); out.save(dst); return out
if __name__=='__main__':
    cut(sys.argv[1], sys.argv[2], *(sys.argv[3:4]))
    print('OK:', sys.argv[2])