Touble Shooting : 영역인식 후처리

조훈·2025년 6월 25일

자율주행 AI트랙 : 주차공간 탐색 프로그램

목록 보기

14/15

AP 결과 수치는 좋게 나왔지만 실제 영상에 적용 시켰을 때
- 영역 경계 부분 불분명
- 영역 곳곳에 빈 부분 발생

1. 해상도 조정

import os
import cv2
import torch
import numpy as np
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor

#  설정
VIDEO_PATH = "YOUR_VID_DIR"
OUTPUT_DIR = "YOUR_OUTPUT_DIR"
MODEL_DIR = "YOUR_MODEL_DIR"
NUM_CLASSES = 3

#  클래스 색상 정의
COLORMAP = np.array([
    [0, 0, 0], [0, 255, 0], [255, 255, 0],
    [0, 0, 255], [255, 0, 0], [255, 165, 0], [255, 0, 255]
], dtype=np.uint8)

#  모델 로드
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_DIR)
processor = SegformerImageProcessor(do_resize=True, size={"height": 1024, "width": 1024}, do_normalize=True)

#  디렉토리 준비
os.makedirs(OUTPUT_DIR, exist_ok=True)

#  비디오 로드
cap = cv2.VideoCapture(VIDEO_PATH)
fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

#  저장용 비디오 설정
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out_video = cv2.VideoWriter(os.path.join(OUTPUT_DIR, "detected_size.mp4"), fourcc, fps, (w, h))

frame_idx = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # OpenCV BGR → PIL RGB
    image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

    # 전처리 및 추론
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
        logits = torch.nn.functional.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
        predicted = logits.argmax(dim=1)[0].cpu().numpy()  # (H, W)

    # 컬러 마스크 → BGR
    color_mask = COLORMAP[predicted]
    color_mask_bgr = cv2.cvtColor(color_mask.astype(np.uint8), cv2.COLOR_RGB2BGR)

    # 원본 프레임과 합성 (반투명)
    overlay = cv2.addWeighted(frame, 0.6, color_mask_bgr, 0.4, 0)

    out_video.write(overlay)
    frame_idx += 1
    if frame_idx % 10 == 0:
        print(f"Processed frame {frame_idx}")

cap.release()
out_video.release()
print("SegFormer segmentation video 저장 완료")

processor - size에서 조절

processor = SegformerImageProcessor(do_resize=True, size={"height": 1024, "width": 1024}, do_normalize=True)

512x512에서 가장 좋은 결과를 보임

2. 시각화 처리 - OpenCV

import os
import cv2
import torch
import numpy as np
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor

# 설정
VIDEO_PATH = "YOUR_VID_DIR"
OUTPUT_DIR = "YOUR_OUTPUT_DIR"
MODEL_DIR = "YOUR_MODEL_DIR"
NUM_CLASSES = 2

# 클래스 색상 정의 (RGB)
COLORMAP = np.array([
    [0, 0, 0], [0, 255, 0], [255, 255, 0],
    [0, 0, 255], [255, 0, 0], [255, 165, 0], [255, 0, 255]
], dtype=np.uint8)

# 모델 로드
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_DIR)
processor = SegformerImageProcessor(do_resize=True, size={"height": 512, "width": 512}, do_normalize=True)

# 디렉토리 생성
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 비디오 설정
cap = cv2.VideoCapture(VIDEO_PATH)
fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out_video = cv2.VideoWriter(os.path.join(OUTPUT_DIR, "detected_overlay_edges.mp4"), fourcc, fps, (w, h))

frame_idx = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # PIL로 변환
    image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        logits = model(**inputs).logits
        logits = torch.nn.functional.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
        predicted = logits.argmax(dim=1)[0].cpu().numpy()  # (H, W)

    # 컬러 마스크 생성 및 BGR 변환
    color_mask = COLORMAP[predicted]
    color_mask_bgr = cv2.cvtColor(color_mask.astype(np.uint8), cv2.COLOR_RGB2BGR)
    color_mask_bgr = cv2.medianBlur(color_mask_bgr, 5)  # 블러 적용

    # 오버레이 생성
    overlay = cv2.addWeighted(frame, 0.6, color_mask_bgr, 0.4, 0)

    # 경계선 추출 (Canny)
    edge_map = cv2.Canny(predicted.astype(np.uint8), 50, 150)
    edges_rgb = cv2.cvtColor(edge_map, cv2.COLOR_GRAY2BGR)
    edges_rgb[np.where((edges_rgb != [0, 0, 0]).all(axis=2))] = [255, 255, 255]  # 흰색 강조

    # 오버레이 + 경계선 합성
    final_frame = cv2.addWeighted(overlay, 1.0, edges_rgb, 0.7, 0)

    out_video.write(final_frame)

    frame_idx += 1
    if frame_idx % 10 == 0:
        print(f"Processed frame {frame_idx}")

cap.release()
out_video.release()
print(" SegFormer + Edge 강조 시각화 영상 저장 완료")

argmax : 픽셀마다 가장 높은 점수를 받은 클래스를 최종 마스크로 추출
medianBlur : 노이즈 제거 및 부드러운 경계 생성
Canny : 경계선 강조

3. CRF 처리

CRF : 조건부 확률 모델의 일종으로, 주어진 입력에 대해 레이블 간의 관계를 고려하여 최적의 라벨링을 결정합니다.

ex ) 하늘은 대부분 잘 분류했지만, 나무에 일부 하늘 라벨이 섞여 있음 → CRF가 이를 보정
python의 pydensecrf 패키지 사용

import os
import cv2
import torch
import numpy as np
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_softmax, create_pairwise_bilateral, create_pairwise_gaussian

# 설정
VIDEO_PATH = "/home/elicer/jhj/output_video.mp4"
OUTPUT_DIR = "/home/elicer/jhj/segformer/detect"
MODEL_DIR = "/home/elicer/jhj/Segformer_Train/Final/result"
NUM_CLASSES = 3

# 클래스 색상 정의 (RGB)
COLORMAP = np.array([
    [0, 0, 0], [0, 255, 0], [255, 255, 0],
    [0, 0, 255], [255, 0, 0], [255, 165, 0], [255, 0, 255]
], dtype=np.uint8)

# 모델 로드
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_DIR)
processor = SegformerImageProcessor(do_resize=True, size={"height": 1024, "width": 1024}, do_normalize=True)

# 디렉토리 준비
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 비디오 로드
cap = cv2.VideoCapture(VIDEO_PATH)
fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# 비디오 저장 설정
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out_video = cv2.VideoWriter(os.path.join(OUTPUT_DIR, "detected_crf.mp4"), fourcc, fps, (w, h))

frame_idx = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # OpenCV BGR → PIL RGB
    image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

    # 전처리 및 추론
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits  # (1, C, H', W')
        logits = torch.nn.functional.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
        probs = torch.softmax(logits[0], dim=0).cpu().numpy()  # (C, H, W)

    # DenseCRF 적용
    crf = dcrf.DenseCRF2D(w, h, NUM_CLASSES)
    unary = unary_from_softmax(probs)
    crf.setUnaryEnergy(unary)

    # 쌍항 조건 추가 (appearance + smoothness)
    image_np = np.asarray(image.resize((w, h)))  # (H, W, 3)
    feats_bilateral = create_pairwise_bilateral(sdims=(80, 80), schan=(13, 13, 13),
                                                 img=image_np, chdim=2)
    feats_gaussian = create_pairwise_gaussian(sdims=(3, 3), shape=(h, w))

    crf.addPairwiseEnergy(feats_gaussian, compat=3)
    crf.addPairwiseEnergy(feats_bilateral, compat=10)

    crf_result = crf.inference(5)
    refined_mask = np.argmax(np.array(crf_result), axis=0).reshape((h, w))  # (H, W)

    # 시각화용 색상 적용
    color_mask = COLORMAP[refined_mask]
    color_mask_bgr = cv2.cvtColor(color_mask.astype(np.uint8), cv2.COLOR_RGB2BGR)

    # 원본 프레임과 합성
    overlay = cv2.addWeighted(frame, 0.6, color_mask_bgr, 0.4, 0)

    out_video.write(overlay)
    frame_idx += 1
    if frame_idx % 10 == 0:
        print(f"Processed frame {frame_idx}")

cap.release()
out_video.release()
print(" SegFormer + CRF 후처리 영상 저장 완료")

오히려 노이즈 발생

해상도를 512x512로 하여 시각화 처리만 적용한 후처리

조훈

이전 포스트

최종 학습 및 평가 (3) : Segformer

다음 포스트

Touble Shooting : 영역인식 후처리

자율주행 AI트랙 : 주차공간 탐색 프로그램

1. 해상도 조정

2. 시각화 처리 - OpenCV

3. CRF 처리

CRF : 조건부 확률 모델의 일종으로, 주어진 입력에 대해 레이블 간의 관계를 고려하여 최적의 라벨링을 결정합니다.

최종 학습 및 평가 (3) : Segformer

최종 학습 및 평가 (4) : 후처리

0개의 댓글