FEAT : 작업3 완료

README.md

@@ -1,10 +1,10 @@
 ---
-title: Segmentation3
-emoji: 📚
-colorFrom: blue
+title: Segment3
+emoji: 🌖
+colorFrom: green
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.2.0
+sdk_version: 3.44.4
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,182 @@
+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
+
+    elif 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(["detection", "segmentation"])
+    
+    return sepia(image_path, selected_option)
+
+# Gr.Dropdown을 사용하여 옵션을 선택할 수 있도록 변경
+dropdown = gr.Dropdown(
+            ["detection", "segmentation"], label="Menu", info="Select One!"
+        )
+
+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()

labels.txt

@@ -0,0 +1,19 @@
+road
+sidewalk
+building
+wall
+fence
+pole
+traffic light
+traffic sign
+vegetation
+terrain
+sky
+person
+rider
+car
+truck
+bus
+train
+motorcycle
+bicycle

requirements.txt

@@ -0,0 +1,7 @@
+torch
+transformers
+tensorflow==2.13.0
+numpy
+Image
+matplotlib
+Pillow