init

app.py

@@ -27,6 +27,8 @@ from data_loader import SalObjDataset
 from model import U2NET # full size version 173.6 MB
 from model import U2NETP # small version u2net 4.7 MB
 
+from modnet import ModNet
+import huggingface_hub
 
 # normalize the predicted SOD probability map
 def normPRED(d):
@@ -58,6 +60,12 @@ def save_output(image_name,pred,d_dir):
     return d_dir+'/'+imidx+'.png'
 
 
+
+modnet_path = huggingface_hub.hf_hub_download('hylee/apdrawing_model',
+                                                      'modnet.onnx',
+                                                      force_filename='modnet.onnx')
+modnet = ModNet(modnet_path)
+
  # --------- 1. get image path and name ---------
 model_name='u2net_portrait'#u2netp
 
@@ -82,6 +90,9 @@ net.eval()
 
 
 def process(im):
+    image = modnet.segment(im.name)
+    Image.fromarray(np.uint8(image)).save(im.name)
+
     img_name_list = glob.glob(im.name)
     print("Number of images: ", len(img_name_list))
     # --------- 2. dataloader ---------

requirements.txt

@@ -3,4 +3,6 @@ scikit-image
 torch
 torchvision
 pillow
-opencv-python-headless
+opencv-python-headless
+onnx==1.8.1
+onnxruntime==1.6.0