Spaces:
Running
Running
File size: 4,105 Bytes
850927f d2aced2 850927f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import gradio as gr
from PIL import Image
import numpy as np
import cv2
import os
import tensorflow as tf
if tf.__version__ >= '2.0':
tf = tf.compat.v1
class ImageMattingPipeline:
def __init__(self, model_dir: str, input_name: str = 'input_image:0', output_name: str = 'output_png:0'):
model_path = os.path.join(model_dir, 'tf_graph.pb')
if not os.path.exists(model_path):
raise FileNotFoundError("Model file not found at {}".format(model_path))
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
self.graph = tf.Graph()
with self.graph.as_default():
self._session = tf.Session(config=config)
with tf.gfile.FastGFile(model_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
self.output = self._session.graph.get_tensor_by_name(output_name)
self.input_name = input_name
def preprocess(self, input_image):
img = np.array(input_image)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img = img.astype(float)
return {'img': img}
def forward(self, input, output_mask=False, alpha_threshold=128):
with self.graph.as_default(), self._session.as_default():
feed_dict = {self.input_name: input['img']}
output_img = self._session.run(self.output, feed_dict=feed_dict)
result = {'output_img': output_img}
if output_mask:
alpha_channel = output_img[:, :, 3]
mask = np.zeros(alpha_channel.shape, dtype=np.uint8)
mask[alpha_channel >= alpha_threshold] = 255
output_img[mask == 0, 3] = 0
result['mask'] = mask
return result
def apply_filters(mask: np.array, closing_kernel: tuple = (5, 5), opening_kernel: tuple = (5, 5),
blur_kernel: tuple = (3, 3), bilateral_params: tuple = (9, 75, 75),
min_area: int = 2000) -> np.array:
mask = mask.astype(np.uint8)
closing_element = np.ones(closing_kernel, np.uint8)
opening_element = np.ones(opening_kernel, np.uint8)
closed_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, closing_element)
opened_mask = cv2.morphologyEx(closed_mask, cv2.MORPH_OPEN, opening_element)
smoothed_mask = cv2.GaussianBlur(opened_mask, blur_kernel, 0)
edge_smoothed_mask = cv2.bilateralFilter(smoothed_mask, *bilateral_params)
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(edge_smoothed_mask, connectivity=8)
large_component_mask = np.zeros_like(edge_smoothed_mask)
for i in range(1, num_labels):
if stats[i, cv2.CC_STAT_AREA] >= min_area:
large_component_mask[labels == i] = 255
return large_component_mask
def matting_interface(input_image, apply_morphology):
input_image = np.array(input_image)
input_image = input_image[:, :, ::-1]
pipeline = ImageMattingPipeline(model_dir='cv_unet_universal-matting')
preprocessed = pipeline.preprocess(input_image)
result = pipeline.forward(preprocessed, output_mask=True)
if apply_morphology:
mask = apply_filters(result['mask'])
else:
mask = result.get('mask', None)
output_img_pil = Image.fromarray(result['output_img'].astype(np.uint8))
mask_pil = Image.fromarray(mask) if mask is not None else None
return output_img_pil, mask_pil
iface = gr.Interface(
fn=matting_interface,
inputs=[
gr.components.Image(type="pil", image_mode="RGB"),
gr.components.Checkbox(label="Apply Morphological Processing for Mask")
],
outputs=[
gr.components.Image(type="pil", label="Matting Result"),
gr.components.Image(type="pil", label="Mask"),
],
title="Image Matting and Mask",
description="Upload an image to get the matting result and mask. "
"Use the checkbox to enable or disable morphological processing on the mask."
)
iface.launch() |