app.py

import os

import gradio as gr

import sys
sys.path.insert(0, 'U-2-Net')

from skimage import io, transform

import numpy as np
from PIL import Image

from utils.face_seg import FaceSeg
import cv2

import requests
import base64
from io import BytesIO

segment = FaceSeg()


def profuai(im_path, out_path):
    r = requests.post(
        'http://nebula.cs.ualberta.ca/predict',
        files={
            'file': open(im_path, 'rb'),
        },
        headers={'Host': 'nebula.cs.ualberta.ca', 'Origin': 'http://nebula.cs.ualberta.ca','Referer':'http://nebula.cs.ualberta.ca/',
                 'X-Requested-With':'XMLHttpRequest',
                 'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.12; rv:51.0) Gecko/20100101 Firefox/51.0'}

    )
    #print(r)
    if (r.status_code == 200):
        data = r.text
        #data:image/png;base64,
        a = data[len("data:image/png;base64,"):]
        missing_padding = 4 - len(a) % 4
        if missing_padding:
            a += '=' * missing_padding

        img = Image.open(BytesIO(base64.urlsafe_b64decode(a)))
        # print(a)
        img.save(out_path, quality=80)
    else:
        raise Exception('error 1001')


def process(im):
    image = cv2.imread(im.name)
    matte = segment.get_mask(image)

    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)
    cv2.imwrite(im.name, foreground)

    profuai(im.name, im.name)

    return Image.open(im.name)
        
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)