File size: 1,930 Bytes
e4c0d76 a45a6f9 e4c0d76 a45a6f9 |
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 |
import gradio as gr
import tensorflow as tf
from PIL import Image
import requests
# ๋ชจ๋ธ ๋ก๋
model = tf.saved_model.load("nvidia_segformer_b5_finetuned_cityscapes_1024")
# ๋ ์ด๋ธ ๋ฐ ์์ ์ ์
label_colors = {
"road": [204, 87, 92],
"sidewalk": [112, 185, 212],
"building": [196, 160, 122],
"wall": [106, 135, 242],
"fence": [91, 192, 222],
"pole": [255, 192, 203],
"traffic_light": [176, 224, 230],
"traffic_sign": [222, 49, 99],
"vegetation": [139, 69, 19],
"terrain": [255, 0, 0],
"sky": [0, 0, 255],
"person": [255, 228, 181],
"rider": [128, 0, 0],
"car": [0, 128, 0],
"truck": [255, 99, 71],
"bus": [0, 255, 0],
"train": [128, 0, 128],
"motorcycle": [255, 255, 0],
"bicycle": [128, 0, 128]
}
# Gradio ์ธํฐํ์ด์ค ์ ์
iface = gr.Interface(
fn=lambda image: predict_segmentation(image, model),
inputs="image",
outputs="image"
)
iface.launch()
# ์ด๋ฏธ์ง ์ธ๊ทธ๋ฉํ
์ด์
ํจ์ ์ ์
def predict_segmentation(image, model):
# ์ด๋ฏธ์ง ๋ณํ
image = Image.fromarray(image.astype('uint8'), 'RGB')
image = image.resize((1024, 1024)) # ๋ชจ๋ธ์ ์
๋ ฅ ํฌ๊ธฐ์ ๋ง๊ฒ ์กฐ์
image_array = tf.keras.preprocessing.image.img_to_array(image)
image_array = tf.expand_dims(image_array, 0)
# ๋ชจ๋ธ ์ถ๋ก
predictions = model(image_array)["output_0"]
# ๋ ์ด๋ธ๋ณ ์์ ๋งคํ
segmented_image = tf.zeros_like(predictions)
for label, color in label_colors.items():
mask = tf.reduce_all(tf.equal(predictions, color), axis=-1, keepdims=True)
for i in range(3):
segmented_image += tf.cast(mask, tf.float32) * tf.constant(color[i], dtype=tf.float32)
# ์ด๋ฏธ์ง ๋ฆฌํด
segmented_image = tf.cast(segmented_image, tf.uint8)
segmented_image = tf.image.resize(segmented_image, [image.height, image.width])
return segmented_image.numpy()
|