gradio/demo/face_segment.py
2020-10-26 15:27:28 -07:00

37 lines
1.2 KiB
Python

import gradio as gr
import os, sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__), "utils"))
from FCN8s_keras import FCN
from PIL import Image
import cv2
import tensorflow as tf
from drive import download_file_from_google_drive
import numpy as np
weights = os.path.join(file_folder, "face_seg_model_weights.h5")
if not os.path.exists(weights):
file_id = "1IerDF2DQqmJWqyvxYZOICJT1eThnG8WR"
download_file_from_google_drive(file_id, weights)
model1 = FCN()
model1.load_weights(weights)
def segment_face(inp):
im = Image.fromarray(np.uint8(inp))
im = im.resize((500, 500))
in_ = np.array(im, dtype=np.float32)
in_ = in_[:, :, ::-1]
in_ -= np.array((104.00698793,116.66876762,122.67891434))
in_ = in_[np.newaxis,:]
out = model1.predict(in_)
out_resized = cv2.resize(np.squeeze(out), (inp.shape[1], inp.shape[0]))
out_resized_clipped = np.clip(out_resized.argmax(axis=2), 0, 1).astype(np.float64)
result = (out_resized_clipped[:, :, np.newaxis] + 0.25)/1.25 * inp.astype(np.float64).astype(np.uint8)
return result / 255
iface = gr.Interface(segment_face, "webcam", "image", capture_session=True)
if __name__ == "__main__":
iface.launch()