init

app.py

@@ -0,0 +1,155 @@
+import os
+
+import gradio as gr
+
+import sys
+sys.path.insert(0, 'U-2-Net')
+
+from skimage import io, transform
+import torch
+import torchvision
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms#, utils
+# import torch.optim as optim
+
+import numpy as np
+from PIL import Image
+import glob
+
+from data_loader import RescaleT
+from data_loader import ToTensor
+from data_loader import ToTensorLab
+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):
+    ma = torch.max(d)
+    mi = torch.min(d)
+
+    dn = (d-mi)/(ma-mi)
+
+    return dn
+def save_output(image_name,pred,d_dir):
+    predict = pred
+    predict = predict.squeeze()
+    predict_np = predict.cpu().data.numpy()
+
+    im = Image.fromarray(predict_np*255).convert('RGB')
+    img_name = image_name.split(os.sep)[-1]
+    image = io.imread(image_name)
+    imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR)
+
+    pb_np = np.array(imo)
+
+    aaa = img_name.split(".")
+    bbb = aaa[0:-1]
+    imidx = bbb[0]
+    for i in range(1,len(bbb)):
+        imidx = imidx + "." + bbb[i]
+
+    imo.save(d_dir+'/'+imidx+'.png')
+    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
+
+
+image_dir = 'portrait_im'
+prediction_dir = 'portrait_results'
+if(not os.path.exists(prediction_dir)):
+    os.mkdir(prediction_dir)
+
+model_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'U-2-Net/saved_models/u2net_portrait/u2net_portrait.pth')
+
+
+# --------- 3. model define ---------
+
+print("...load U2NET---173.6 MB")
+net = U2NET(3,1)
+
+net.load_state_dict(torch.load(model_dir, map_location='cpu'))
+# if torch.cuda.is_available():
+#     net.cuda()
+net.eval()
+
+
+def process(im):
+    image = modnet.segment(im.name)
+    im_path = os.path.abspath(os.path.basename(im.name))
+    Image.fromarray(np.uint8(image)).save(im_path)
+
+    img_name_list = [im_path]
+    print("Number of images: ", len(img_name_list))
+    # --------- 2. dataloader ---------
+    # 1. dataloader
+    test_salobj_dataset = SalObjDataset(img_name_list=img_name_list,
+                                        lbl_name_list=[],
+                                        transform=transforms.Compose([RescaleT(512),
+                                                                      ToTensorLab(flag=0)])
+                                        )
+    test_salobj_dataloader = DataLoader(test_salobj_dataset,
+                                        batch_size=1,
+                                        shuffle=False,
+                                        num_workers=1)
+
+    results = []
+    # --------- 4. inference for each image ---------
+    for i_test, data_test in enumerate(test_salobj_dataloader):
+
+        print("inferencing:", img_name_list[i_test].split(os.sep)[-1])
+
+        inputs_test = data_test['image']
+        inputs_test = inputs_test.type(torch.FloatTensor)
+
+        # if torch.cuda.is_available():
+        #     inputs_test = Variable(inputs_test.cuda())
+        # else:
+        inputs_test = Variable(inputs_test)
+
+        d1, d2, d3, d4, d5, d6, d7 = net(inputs_test)
+
+        # normalization
+        pred = 1.0 - d1[:, 0, :, :]
+        pred = normPRED(pred)
+
+        # save results to test_results folder
+        results.append(save_output(img_name_list[i_test], pred, prediction_dir))
+
+        del d1, d2, d3, d4, d5, d6, d7
+
+    print(results)
+
+    return Image.open(results[0])
+        
+title = "U-2-Net"
+description = "Gradio demo for U-2-Net, https://github.com/xuebinqin/U-2-Net"
+article = ""
+
+gr.Interface(
+    process, 
+    [gr.inputs.Image(type="file", label="Input")
+], 
+    [gr.outputs.Image(type="pil", label="Output")],
+    title=title,
+    description=description,
+    article=article,
+    examples=[],
+    allow_flagging=False,
+    allow_screenshot=False
+    ).launch(enable_queue=True,cache_examples=True)

modnet.py

@@ -0,0 +1,94 @@
+import os
+import cv2
+import argparse
+import numpy as np
+from PIL import Image
+
+import onnx
+import onnxruntime
+
+
+class ModNet:
+
+    def __init__(self, model_path):
+        # Initialize session and get prediction
+        self.session = onnxruntime.InferenceSession(model_path, None)
+
+    # Get x_scale_factor & y_scale_factor to resize image
+    def get_scale_factor(self, im_h, im_w, ref_size):
+
+        if max(im_h, im_w) < ref_size or min(im_h, im_w) > ref_size:
+            if im_w >= im_h:
+                im_rh = ref_size
+                im_rw = int(im_w / im_h * ref_size)
+            elif im_w < im_h:
+                im_rw = ref_size
+                im_rh = int(im_h / im_w * ref_size)
+        else:
+            im_rh = im_h
+            im_rw = im_w
+
+        im_rw = im_rw - im_rw % 32
+        im_rh = im_rh - im_rh % 32
+
+        x_scale_factor = im_rw / im_w
+        y_scale_factor = im_rh / im_h
+
+        return x_scale_factor, y_scale_factor
+
+    def segment(self, image_path):
+        ref_size = 512
+        ##############################################
+        #  Main Inference part
+        ##############################################
+
+        # read image
+        im = cv2.imread(image_path)
+        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
+
+        # unify image channels to 3
+        if len(im.shape) == 2:
+            im = im[:, :, None]
+        if im.shape[2] == 1:
+            im = np.repeat(im, 3, axis=2)
+        elif im.shape[2] == 4:
+            im = im[:, :, 0:3]
+
+        # normalize values to scale it between -1 to 1
+        im = (im - 127.5) / 127.5
+
+        im_h, im_w, im_c = im.shape
+        x, y = self.get_scale_factor(im_h, im_w, ref_size)
+
+        # resize image
+        im = cv2.resize(im, None, fx=x, fy=y, interpolation=cv2.INTER_AREA)
+
+        # prepare input shape
+        im = np.transpose(im)
+        im = np.swapaxes(im, 1, 2)
+        im = np.expand_dims(im, axis=0).astype('float32')
+
+        input_name = self.session.get_inputs()[0].name
+        output_name = self.session.get_outputs()[0].name
+        result = self.session.run([output_name], {input_name: im})
+
+        # refine matte
+        matte = (np.squeeze(result[0]) * 255).astype('uint8')
+        matte = cv2.resize(matte, dsize=(im_w, im_h), interpolation=cv2.INTER_AREA)
+
+        # obtain predicted foreground
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        if len(image.shape) == 2:
+            image = image[:, :, None]
+        if image.shape[2] == 1:
+            image = np.repeat(image, 3, axis=2)
+        elif image.shape[2] == 4:
+            image = image[:, :, 0:3]
+        matte = np.repeat(np.asarray(matte)[:, :, None], 3, axis=2) / 255
+        foreground = image * matte + np.full(image.shape, 255) * (1 - matte)
+
+        return foreground
+
+

requirements.txt

@@ -0,0 +1,8 @@
+numpy
+scikit-image
+torch
+torchvision
+pillow
+opencv-python-headless
+onnx==1.8.1
+onnxruntime==1.6.0