segmentation3 / app.py
SeoJunn's picture
Refactor : detection ์ฝ”๋“œ ์ฃผ์„
bb78b14
raw
history blame
5.95 kB
import gradio as gr
from PIL import Image
from matplotlib import gridspec
import matplotlib.pyplot as plt
import numpy as np
from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
import tensorflow as tf
from PIL import ImageDraw
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# image segmentation ๋ชจ๋ธ
feature_extractor = SegformerFeatureExtractor.from_pretrained(
"nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)
model_segmentation = TFSegformerForSemanticSegmentation.from_pretrained(
"nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)
# image detection ๋ชจ๋ธ
# processor_detection = DetrImageProcessor.from_pretrained(
# "facebook/detr-resnet-50", revision="no_timm"
# )
# model_detection = DetrForObjectDetection.from_pretrained(
# "facebook/detr-resnet-50", revision="no_timm"
# )
def ade_palette():
"""ADE20K ํŒ”๋ ˆํŠธ: ๊ฐ ํด๋ž˜์Šค๋ฅผ RGB ๊ฐ’์— ๋งคํ•‘ํ•ด์ฃผ๋Š” ํ•จ์ˆ˜์ž…๋‹ˆ๋‹ค."""
return [
[204, 87, 92],
[112, 185, 212],
[45, 189, 106],
[234, 123, 67],
[78, 56, 123],
[210, 32, 89],
[90, 180, 56],
[155, 102, 200],
[33, 147, 176],
[255, 183, 76],
[67, 123, 89],
[190, 60, 45],
[134, 112, 200],
[56, 45, 189],
[200, 56, 123],
[87, 92, 204],
[120, 56, 123],
[45, 78, 123],
[45, 123, 67],
]
labels_list = []
with open(r"labels.txt", "r") as fp:
for line in fp:
labels_list.append(line[:-1])
colormap = np.asarray(ade_palette())
def label_to_color_image(label):
"""๋ผ๋ฒจ์„ ์ปฌ๋Ÿฌ ์ด๋ฏธ์ง€๋กœ ๋ณ€ํ™˜ํ•ด์ฃผ๋Š” ํ•จ์ˆ˜์ž…๋‹ˆ๋‹ค."""
if label.ndim != 2:
raise ValueError("2์ฐจ์› ์ž…๋ ฅ ๋ผ๋ฒจ์„ ๊ธฐ๋Œ€ํ•ฉ๋‹ˆ๋‹ค.")
if np.max(label) >= len(colormap):
raise ValueError("๋ผ๋ฒจ ๊ฐ’์ด ๋„ˆ๋ฌด ํฝ๋‹ˆ๋‹ค.")
return colormap[label]
def draw_plot(pred_img, seg):
"""์ด๋ฏธ์ง€์™€ ์„ธ๊ทธ๋ฉ˜ํ…Œ์ด์…˜ ๊ฒฐ๊ณผ๋ฅผ floating ํ•˜๋Š” ํ•จ์ˆ˜์ž…๋‹ˆ๋‹ค."""
fig = plt.figure(figsize=(20, 15))
grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
plt.subplot(grid_spec[0])
plt.imshow(pred_img)
plt.axis("off")
LABEL_NAMES = np.asarray(labels_list)
FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
unique_labels = np.unique(seg.numpy().astype("uint8"))
ax = plt.subplot(grid_spec[1])
plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
ax.yaxis.tick_right()
plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
plt.xticks([], [])
ax.tick_params(width=0.0, labelsize=25)
return fig
def sepia(inputs, button_text):
"""๊ฐ์ฒด ๊ฒ€์ถœ ๋˜๋Š” ์„ธ๊ทธ๋ฉ˜ํ…Œ์ด์…˜์„ ์ˆ˜ํ–‰ํ•˜๊ณ  ๊ฒฐ๊ณผ๋ฅผ ๋ฐ˜ํ™˜ํ•˜๋Š” ํ•จ์ˆ˜์ž…๋‹ˆ๋‹ค."""
input_img = Image.fromarray(inputs)
# if button_text == "detection":
# inputs_detection = processor_detection(images=input_img, return_tensors="pt")
# outputs_detection = model_detection(**inputs_detection)
# target_sizes = torch.tensor([input_img.size[::-1]])
# results_detection = processor_detection.post_process_object_detection(
# outputs_detection, target_sizes=target_sizes, threshold=0.9
# )[0]
# draw = ImageDraw.Draw(input_img)
# for score, label, box in zip(
# results_detection["scores"],
# results_detection["labels"],
# results_detection["boxes"],
# ):
# box = [round(i, 2) for i in box.tolist()]
# label_name = model_detection.config.id2label[label.item()]
# print(
# f"Detected {label_name} with confidence "
# f"{round(score.item(), 3)} at location {box}"
# )
# draw.rectangle(box, outline="red", width=3)
# draw.text((box[0], box[1]), label_name, fill="red", font=None)
# fig = plt.figure(figsize=(20, 15))
# plt.imshow(input_img)
# plt.axis("off")
# return fig
if button_text == "segmentation":
inputs_segmentation = feature_extractor(images=input_img, return_tensors="tf")
outputs_segmentation = model_segmentation(**inputs_segmentation)
logits_segmentation = outputs_segmentation.logits
logits_segmentation = tf.transpose(logits_segmentation, [0, 2, 3, 1])
logits_segmentation = tf.image.resize(logits_segmentation, input_img.size[::-1])
seg = tf.math.argmax(logits_segmentation, axis=-1)[0]
color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
for label, color in enumerate(colormap):
color_seg[seg.numpy() == label, :] = color
pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
pred_img = pred_img.astype(np.uint8)
fig = draw_plot(pred_img, seg)
return fig
return "Please select 'detection' or 'segmentation'."
def on_button_click(inputs):
"""๋ฒ„ํŠผ ํด๋ฆญ ์ด๋ฒคํŠธ ํ•ธ๋“ค๋Ÿฌ"""
image_path, selected_option = inputs
if selected_option == "dropout":
# 'dropout'์ด๋ฉด ๋‘ ๊ฐ€์ง€ ์ค‘์— ํ•˜๋‚˜๋ฅผ ๋žœ๋ค์œผ๋กœ ์„ ํƒ
selected_option = np.random.choice(["segmentation"])
return sepia(image_path, selected_option)
# Gr.Dropdown์„ ์‚ฌ์šฉํ•˜์—ฌ ์˜ต์…˜์„ ์„ ํƒํ•  ์ˆ˜ ์žˆ๋„๋ก ๋ณ€๊ฒฝ
dropdown = gr.Dropdown(
["segmentation"], label="Menu", info="Chose Segmentation!"
)
demo = gr.Interface(
fn=sepia,
inputs=[gr.Image(shape=(400, 600)), dropdown],
outputs=["plot"],
examples=[
["01.jpg", "Click me"],
["02.jpeg", "Click me"],
["03.jpeg", "Click me"],
["04.jpeg", "Click me"],
],
allow_flagging="never",
)
demo.launch()