토목공학과

# ============================================================

# Segmentation Practice for Google Colab

# Original Image + VIA CSV

# Polygon / Polyline / Rect / Circle / Ellipse 지원

# ============================================================

import cv2

import json

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

from PIL import Image

from PIL import ImageDraw

from google.colab import files

# ============================================================

# 1. 파일 업로드

# ============================================================

print("원본 이미지와 VIA CSV 파일을 업로드하세요.")

uploaded = files.upload()

image_file = None

csv_file = None

for f in uploaded.keys():

    if f.lower().endswith(

        (".jpg",".jpeg",".png")

    ):

        image_file = f

    elif f.lower().endswith(".csv"):

        csv_file = f

print("\n원본 이미지 :", image_file)

print("CSV 파일 :", csv_file)

if image_file is None:

    raise Exception("이미지 파일 없음")

if csv_file is None:

    raise Exception("CSV 파일 없음")

# ============================================================

# 2. 이미지 읽기

# ============================================================

image = Image.open(image_file).convert("RGB")

image_np = np.array(image)

height, width = image_np.shape[:2]

print("\nImage Size :", width, "x", height)

# ============================================================

# 3. VIA CSV → Ground Truth Mask

# ============================================================

def create_mask_from_via_csv(

    csv_path,

    image_size

):

    width, height = image_size

    mask = Image.new(

        "L",

        (width,height),

        0

    )

    draw = ImageDraw.Draw(mask)

    df = pd.read_csv(csv_path)

    for _, row in df.iterrows():

        shape_str = row.get(

            "region_shape_attributes",

            ""

        )

        if pd.isna(shape_str):

            continue

        if shape_str == "{}":

            continue

        try:

            shape = json.loads(shape_str)

        except:

            continue

        shape_name = shape.get(

            "name",

            ""

        )

        # ---------------------------------

        # polygon

        # ---------------------------------

        if shape_name == "polygon":

            xs = shape.get(

                "all_points_x",

                []

            )

            ys = shape.get(

                "all_points_y",

                []

            )

            pts = list(zip(xs,ys))

            if len(pts) >= 3:

                draw.polygon(

                    pts,

                    outline=255,

                    fill=255

                )

        # ---------------------------------

        # polyline

        # ---------------------------------

        elif shape_name == "polyline":

            xs = shape.get(

                "all_points_x",

                []

            )

            ys = shape.get(

                "all_points_y",

                []

            )

            pts = list(zip(xs,ys))

            if len(pts) >= 3:

                draw.polygon(

                    pts,

                    outline=255,

                    fill=255

                )

        # ---------------------------------

        # rect

        # ---------------------------------

        elif shape_name == "rect":

            x = shape.get("x",0)

            y = shape.get("y",0)

            w = shape.get("width",0)

            h = shape.get("height",0)

            draw.rectangle(

                [x,y,x+w,y+h],

                outline=255,

                fill=255

            )

        # ---------------------------------

        # circle

        # ---------------------------------

        elif shape_name == "circle":

            cx = shape.get("cx",0)

            cy = shape.get("cy",0)

            r  = shape.get("r",0)

            draw.ellipse(

                [

                    cx-r,

                    cy-r,

                    cx+r,

                    cy+r

                ],

                outline=255,

                fill=255

            )

        # ---------------------------------

        # ellipse

        # ---------------------------------

        elif shape_name == "ellipse":

            cx = shape.get("cx",0)

            cy = shape.get("cy",0)

            rx = shape.get("rx",0)

            ry = shape.get("ry",0)

            draw.ellipse(

                [

                    cx-rx,

                    cy-ry,

                    cx+rx,

                    cy+ry

                ],

                outline=255,

                fill=255

            )

    return np.array(mask)

# ============================================================

# GT 생성

# ============================================================

gt_mask = create_mask_from_via_csv(

    csv_file,

    (width,height)

)

# ============================================================

# 4. Prediction 생성 (실습용)

# ============================================================

kernel = np.ones(

    (15,15),

    np.uint8

)

pred_mask = cv2.GaussianBlur(

    gt_mask,

    (31,31),

    0

)

_, pred_mask = cv2.threshold(

    pred_mask,

    80,

    255,

    cv2.THRESH_BINARY

)

pred_mask = cv2.erode(

    pred_mask,

    kernel,

    iterations=1

)

pred_mask = cv2.dilate(

    pred_mask,

    kernel,

    iterations=1

)

# ============================================================

# 5. Overlay 생성

# ============================================================

gt_bin = gt_mask > 0

pred_bin = pred_mask > 0

intersection = gt_bin & pred_bin

difference = gt_bin ^ pred_bin

overlay = image_np.copy()

overlay[intersection] = [0,255,0]

overlay[difference] = [255,0,0]

overlay_blend = cv2.addWeighted(

    image_np,

    0.6,

    overlay,

    0.4,

    0

)

# ============================================================

# 6. IoU Dice

# ============================================================

def calc_iou_dice(

    gt,

    pred

):

    gt = gt > 0

    pred = pred > 0

    inter = np.logical_and(

        gt,

        pred

    ).sum()

    union = np.logical_or(

        gt,

        pred

    ).sum()

    iou = inter / (union + 1e-7)

    dice = (

        2*inter

    ) / (

        gt.sum()

        + pred.sum()

        + 1e-7

    )

    return iou,dice

iou,dice = calc_iou_dice(

    gt_mask,

    pred_mask

)

print("\n======================")

print("IoU :", round(iou,4))

print("Dice:", round(dice,4))

print("======================")

# ============================================================

# 7. 결과 표시

# ============================================================

fig, ax = plt.subplots(

    1,

    4,

    figsize=(20,5)

)

ax[0].imshow(image_np)

ax[0].set_title("Original")

ax[0].axis("off")

ax[1].imshow(

    gt_mask,

    cmap="gray"

)

ax[1].set_title("Ground Truth")

ax[1].axis("off")

ax[2].imshow(

    pred_mask,

    cmap="gray"

)

ax[2].set_title("Prediction")

ax[2].axis("off")

ax[3].imshow(

    overlay_blend

)

ax[3].set_title(

    f"Overlay\nIoU={iou:.3f} Dice={dice:.3f}"

)

ax[3].axis("off")

plt.tight_layout()

plt.show()

# ============================================================

# 8. 저장

# ============================================================

Image.fromarray(

    gt_mask

).save(

    "ground_truth_mask.png"

)

Image.fromarray(

    pred_mask

).save(

    "prediction_mask.png"

)

Image.fromarray(

    overlay_blend

).save(

    "overlay_result.png"

)

print("\n저장 완료")

# ============================================================

# 9. 다운로드

# ============================================================

files.download(

    "ground_truth_mask.png"

)

files.download(

    "prediction_mask.png"

)

files.download(

    "overlay_result.png"

)