Spaces:
Sleeping
Sleeping
File size: 1,702 Bytes
9b9f78f 6446bac 9b9f78f |
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 |
import gradio as gr
import numpy as np
import tensorflow as tf
import cv2
import os
# Suppress TensorFlow logging
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
# Constants
IMAGE_H = 512
IMAGE_W = 512
NUM_CLASSES = 11
# RGB color codes for each class
RGB_CODES = [
[0, 0, 0], [0, 153, 255], [102, 255, 153], [0, 204, 153],
[255, 255, 102], [255, 255, 204], [255, 153, 0], [255, 102, 255],
[102, 0, 51], [255, 204, 255], [255, 0, 102]
]
# Load the trained model
model = tf.keras.models.load_model("ten_epoch_model.h5")
# Function to convert grayscale mask to RGB mask
def grayscale_to_rgb(mask, rgb_codes):
h, w = mask.shape[0], mask.shape[1]
mask = mask.astype(np.int32)
output = [rgb_codes[pixel] for pixel in mask.flatten()]
return np.reshape(output, (h, w, 3)).astype(np.uint8)
# Gradio inference function
def segment_face(image):
# Resize and normalize input image
image_resized = cv2.resize(image, (IMAGE_W, IMAGE_H))
image_input = image_resized / 255.0
image_input = np.expand_dims(image_input, axis=0).astype(np.float32)
# Predict the mask
pred = model.predict(image_input, verbose=0)[0]
pred_mask = np.argmax(pred, axis=-1).astype(np.uint8)
# Convert predicted mask to RGB
rgb_mask = grayscale_to_rgb(pred_mask, RGB_CODES)
return rgb_mask
# Launch Gradio app
iface = gr.Interface(
fn=segment_face,
inputs=gr.Image(type="numpy", label="Upload Face Image"),
outputs=gr.Image(type="numpy", label="Segmentation Mask"),
title="Face Segmentation",
description="Upload a face image to get a segmentation mask with different facial components marked in color."
)
if __name__ == "__main__":
iface.launch()
|