[인공지능 프로젝트] 저시력자를 위한 실내 근거리 물체 탐지 및 알림 시스템 - 6편 (통합 시스템 구축)

이전 글에서는 "저시력자를 위한 실내 근거리 물체 탐지 및 알림 시스템" 프로젝트에서 깊이 검출 모델을 어떻게 학습시키는 지와 결과 해석에 대해서 다루었습니다.

 

• https://whitecode2718.tistory.com/163

[인공지능 프로젝트] 저시력자를 위한 실내 근거리 물체 탐지 및 알림 시스템 - 4편 (YOLOv11n 모델

 

본 글에서는 프로젝트에서 진행한 객체 검출 모델과 깊이 추정 모델을 통합한 아키텍쳐와 코드에 대해 간단하게 소개하겠습니다.

 

📌 깃허브 링크: https://github.com/lko9911/Artificial-Intelligence-Project

GitHub - lko9911/Artificial-Intelligence-Project: 강원대학교 인공지능 수업_프로젝트

 

 

📌 결과 추론: https://github.com/lko9911/Artificial-Intelligence-Project/blob/main/Prototype/result_conbine.py

Artificial-Intelligence-Project/Prototype/result_conbine.py at main · lko9911/Artificial-Intelligence-Project

 

 

📌 실시간 : https://github.com/lko9911/Artificial-Intelligence-Project/blob/main/Prototype/prototype_alarm_new.py

Artificial-Intelligence-Project/Prototype/prototype_alarm_new.py at main · lko9911/Artificial-Intelligence-Project

 

---

 

1. 모델 통합 파이프라인

 

지금까지 구현한 검출 기능과 깊이추정 기능을 합하여 저시력자를 위한 실내 근거리 물체 탐지 및 알림 시스템을 구현하겠습니다.

 

기본적인 파이프 라인은 위와 같습니다. 영상을 GUI (시스템)에 입력하면 YOLOv11n 모델이 객체를 검출하고, DeepLabV+ 모델을 통해 영상의 깊이맵을 추론합니다.

 

이때 결과는 YOLO 모델이 검출한 대상의 바운딩 박스내의 깊이맵을 계산하여 해당 객체의 위치 좌표를 반환하게 됩니다. 

 

 

위치 거리를 알게된 대상에 대한 시스템이 프로젝트의 주 목적이기 때문에, 대상의 거리가 특정 거리 이내에 온다면 YOLO + OpenCV를 통해 바운딩 박스의 색상을 변화 시키면서 알람 기능을 추가할수 있습니다.

 

이때 색맹 환자의 경우 색상 구분이 힘들수 있기 때문에 이것또한 사용자가 임의로 색상조합을 선택할수 있도록 하였습니다.

 

2. 기능 동시 비교

import os
import random
import torch
import numpy as np
import cv2 
import matplotlib.pyplot as plt
from ultralytics import YOLO 
from torchvision import transforms 
from PIL import Image 

# Depth Estimation 모델 관련 파일 (사용자 환경에 맞게 로드)
# 이 파일들은 동일한 디렉토리에 있어야 합니다.
from dataset import DepthDataset 
from model import DeepLabv3Plus_Depth 

# ==========================
# 1. 기본 설정 및 경로
# ==========================
device = "cuda" if torch.cuda.is_available() else "cpu"
eps = 1e-6

# 단일 이미지 경로
source = "Depth/SUNRGBD/kv1/NYUdata/NYU0005/image/NYU0005.jpg" 

# Depth 모델의 체크포인트 경로
checkpoint_path = "Depth/checkpoints_depth/best.pth"

# DepthDataset의 img_size와 일치하는 크기 (width, height)
TARGET_SIZE = (512, 256) 
TARGET_WIDTH, TARGET_HEIGHT = TARGET_SIZE 


# ==========================
# 2. YOLO 모델 설정 및 추론
# ==========================
print(f"[{os.path.basename(__file__)}] Starting analysis...")
yolo_model = YOLO("YOLO/model/best.pt")

# YOLO 추론 실행
yolo_results = yolo_model.predict(source, save=False, conf=0.5, show=False)

if yolo_results:
    yolo_im_array_bgr_original = yolo_results[0].plot() 
    
    # YOLO 결과물을 TARGET_SIZE로 리사이즈하여 크기 통일
    yolo_im_array_bgr_resized = cv2.resize(yolo_im_array_bgr_original, TARGET_SIZE, interpolation=cv2.INTER_AREA)
    yolo_im_array_rgb = cv2.cvtColor(yolo_im_array_bgr_resized, cv2.COLOR_BGR2RGB)
    print(f"[DEBUG] YOLO Output Shape: {yolo_im_array_rgb.shape}")
else:
    print("[WARNING] YOLO Prediction failed or no object detected.")
    yolo_im_array_rgb = None


# ==========================
# 3. Depth Estimation 모델 설정 및 추론
# ==========================
depth_model = DeepLabv3Plus_Depth(output_channels=1).to(device)

if os.path.exists(checkpoint_path):
    checkpoint = torch.load(checkpoint_path, map_location=device)
    depth_model.load_state_dict(checkpoint["model_state"])
    print(f"[INFO] Loaded Depth checkpoint: {checkpoint_path}")
else:
    print("[WARNING] Depth checkpoint not found. Using untrained model.")

depth_model.eval()

# --- 입력 이미지 및 GT Depth 맵 로드 및 전처리 ---

try:
    # 1. Depth GT 맵 로드 및 전처리 (DepthDataset의 로직과 동일: Normalize 없음)
    depth_gt_source = source.replace("image", "depth").replace(".jpg", ".png")
    depth_gt_map_raw = cv2.imread(depth_gt_source, cv2.IMREAD_UNCHANGED).astype("float32") / 1000.0
    depth_gt_map_raw[depth_gt_map_raw <= 0] = 0.01

    # 2. RGB 입력 이미지 로드 및 전처리
    image_raw = Image.open(source).convert('RGB')
    
    # DepthDataset의 전처리: Resize(PIL) -> ToTensor() 만 적용
    transform = transforms.Compose([
        transforms.Resize(TARGET_SIZE), 
        transforms.ToTensor(),
    ])
    
    image_tensor = transform(image_raw)
    
    # GT Depth 맵 전처리 
    depth_gt_resized = cv2.resize(depth_gt_map_raw, TARGET_SIZE) 
    depth_gt_tensor = torch.from_numpy(depth_gt_resized).unsqueeze(0)
    
    # 추론 준비
    image_batch = image_tensor.unsqueeze(0).to(device)
    depth_gt = depth_gt_tensor.squeeze().numpy()
    
    # 시각화용 이미지 
    image_np_display = np.array(image_raw.resize(TARGET_SIZE)) / 255.0
    
    # 3. Depth 추론
    with torch.no_grad():
        depth_pred_tensor = depth_model(image_batch)
        depth_pred_raw = depth_pred_tensor.squeeze().cpu().numpy()
        
        # ★★★ 문제 해결: Height x Width 순서로 맞추기 위해 전치(Transpose) 적용 ★★★
        depth_pred = depth_pred_raw.T 
        print(f"[DEBUG] Predicted Depth Shape (AFTER TRANSPOSE): {depth_pred.shape}")


    # 4. 시각화를 위한 깊이 정규화 
    def normalize_depth(depth_map):
        mask = depth_map > 0
        min_d = np.min(depth_map[mask]) if np.any(mask) else 0
        max_d = np.max(depth_map[mask]) if np.any(mask) else 1
        return (depth_map - min_d) / (max_d - min_d + 1e-6)

    pred_vis = normalize_depth(depth_pred) 
    gt_vis = normalize_depth(depth_gt) 
    
    print(f"[DEBUG] Input Image Display Shape: {image_np_display.shape}")
    print(f"[DEBUG] GT Depth Shape: {gt_vis.shape}")
    
except Exception as e:
    print(f"[ERROR] Depth 추론 또는 로드 중 오류 발생: {e}")
    is_depth_ok = False
else:
    is_depth_ok = True


# ==========================
# 4. 통합 결과 시각화
# ==========================
if yolo_im_array_rgb is not None and is_depth_ok:
    # 4개 모두 출력
    fig, axs = plt.subplots(1, 4, figsize=(18, 4)) 
    
    # 모든 이미지에 aspect='equal'을 적용하여 왜곡 방지
    
    # (1) 입력 이미지 
    axs[0].imshow(image_np_display, aspect='equal') 
    axs[0].set_title("Input Image", fontsize=12)
    
    # (2) YOLO 예측 결과 
    axs[1].imshow(yolo_im_array_rgb, aspect='equal') 
    axs[1].set_title(f"YOLO Detection (Conf > {0.5})", fontsize=12)

    # (3) Ground Truth Depth 
    axs[2].imshow(gt_vis, cmap='plasma', aspect='equal') 
    axs[2].set_title("Ground Truth Depth", fontsize=12)
    
    # (4) Predicted Depth (이제 순서가 올바름)
    axs[3].imshow(pred_vis, cmap='plasma', aspect='equal') 
    axs[3].set_title("Predicted Depth", fontsize=12)
    

    # 축 정보 제거
    for ax in axs:
        ax.axis('off')

    plt.suptitle(f"Integrated Analysis for: {os.path.basename(source)} (Resized to {TARGET_WIDTH}x{TARGET_HEIGHT})", fontsize=14)
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.show()

else:
    print("\n[FAIL] 시각화 실패: YOLO 또는 Depth 모델 추론에 필요한 모든 요소가 준비되지 않았습니다.")

 

 

테스트 결과 (YOLO / Depth 개선 전에 테스트한 이미지 입니다.)

 

 

---

 

3. 통합 코드 (Cam 버전)

import os
import random
import torch
import numpy as np
import cv2
from ultralytics import YOLO
from torchvision import transforms
from PIL import Image
import sys

from dataset import DepthDataset
from model import DeepLabv3Plus_Depth

# =============================================================
# 기본 설정
# =============================================================
device = "cuda" if torch.cuda.is_available() else "cpu"
eps = 1e-6

checkpoint_path = "checkpoints_depth/best2.pth"

TARGET_SIZE = (512, 256)
TARGET_WIDTH, TARGET_HEIGHT = TARGET_SIZE

# =============================================================
# 거리 기반 색상
# =============================================================
def get_distance_color(distance, warning_threshold=7):
    if distance < warning_threshold:
        return (255, 0, 0)   # 빨강
    else:
        return (0, 255, 0)   # 초록

# =============================================================
# 카메라 내부 파라미터 (NYU 기준)
# =============================================================
FX_ORIGINAL = 518.857901
FY_ORIGINAL = 519.469611
CX_ORIGINAL = 284.582449
CY_ORIGINAL = 208.736166
ORIGINAL_WIDTH = 640.0
ORIGINAL_HEIGHT = 480.0

SCALE_X_K = TARGET_WIDTH / ORIGINAL_WIDTH
SCALE_Y_K = TARGET_HEIGHT / ORIGINAL_HEIGHT

FX = FX_ORIGINAL * SCALE_X_K
FY = FY_ORIGINAL * SCALE_Y_K
CX = CX_ORIGINAL * SCALE_X_K
CY = CY_ORIGINAL * SCALE_Y_K

K_MATRIX = np.array([
    [FX, 0, CX],
    [0, FY, CY],
    [0, 0, 1]
])
K_inv = np.linalg.inv(K_MATRIX)

# =============================================================
# YOLO 모델 로딩
# =============================================================
yolo_model = YOLO("YOLO/model/best.pt")

# =============================================================
# Depth 모델 로딩
# =============================================================
depth_model = DeepLabv3Plus_Depth(output_channels=1).to(device)

if not os.path.exists(checkpoint_path):
    print("[ERROR] 체크포인트 없음")
    sys.exit(0)

checkpoint = torch.load(checkpoint_path, map_location=device)
depth_model.load_state_dict(checkpoint["model_state"])
depth_model.eval()

transform = transforms.Compose([
    transforms.Resize(TARGET_SIZE),
    transforms.ToTensor()
])

# =============================================================
# 자동 웹캠 탐색
# =============================================================
def find_working_camera(max_index=5):
    for i in range(max_index):
        cap = cv2.VideoCapture(i, cv2.CAP_DSHOW)
        if cap.isOpened():
            print(f"[OK] 카메라 열림: 인덱스 {i}")
            return cap
        cap.release()
    return None

cap = find_working_camera()
if cap is None:
    print("[ERROR] 사용 가능한 웹캠을 찾을 수 없습니다.")
    sys.exit(0)

# 웹캠 해상도 설정
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

# =============================================================
# 창 크기 고정 (1280 × 720)
# =============================================================
WINDOW_W, WINDOW_H = 1280, 720
cv2.namedWindow("YOLO + Depth 3D", cv2.WINDOW_NORMAL)
cv2.resizeWindow("YOLO + Depth 3D", WINDOW_W, WINDOW_H)

cv2.namedWindow("Depth Map", cv2.WINDOW_NORMAL)
cv2.resizeWindow("Depth Map", WINDOW_W, WINDOW_H)

print("\n===== 실시간 추론 시작 (종료: Q) =====")

# =============================================================
# 실시간 루프
# =============================================================
while True:
    ret, frame = cap.read()
    if not ret:
        print("[ERROR] 웹캠 프레임 읽기 실패")
        break

    # PIL 변환
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    image_pil = Image.fromarray(frame_rgb)

    # YOLO 추론
    yolo_results = yolo_model.predict(frame_rgb, save=False, conf=0.5, verbose=False)

    # 리사이즈
    resized_np = cv2.resize(frame_rgb, TARGET_SIZE)

    detected_objects = []

    if yolo_results and len(yolo_results[0].boxes) > 0:
        original_h, original_w = frame_rgb.shape[:2]
        scale_x = TARGET_WIDTH / original_w
        scale_y = TARGET_HEIGHT / original_h

        for box in yolo_results[0].boxes:
            x1o, y1o, x2o, y2o = box.xyxy[0].cpu().numpy().astype(int)
            class_id = int(box.cls[0].item())
            class_name = yolo_results[0].names[class_id]

            x1 = int(x1o * scale_x)
            y1 = int(y1o * scale_y)
            x2 = int(x2o * scale_x)
            y2 = int(y2o * scale_y)

            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2

            detected_objects.append({
                'class': class_name,
                'center_2d': (center_x, center_y),
                'bbox_xyxy': (x1, y1, x2, y2),
                'center_3d': None
            })

    # Depth 예측
    depth_input = transform(image_pil).unsqueeze(0).to(device)

    with torch.no_grad():
        depth_pred = depth_model(depth_input).squeeze().cpu().numpy().T

    depth_norm = cv2.normalize(depth_pred, None, 0, 255, cv2.NORM_MINMAX)
    depth_norm = depth_norm.astype(np.uint8)
    depth_color = cv2.applyColorMap(depth_norm, cv2.COLORMAP_MAGMA)

    # ======= 3D 계산 & 화면 표시 =======
    display_img = resized_np.copy()

    for obj in detected_objects:
        cx, cy = obj['center_2d']
        cx = np.clip(cx, 0, TARGET_WIDTH - 1)
        cy = np.clip(cy, 0, TARGET_HEIGHT - 1)

        Z = depth_pred[cy, cx]

        pixel = np.array([cx, cy, 1])
        X, Y, Z_cam = (Z * (K_inv @ pixel))

        obj['center_3d'] = (X, Y, Z_cam)

        color = get_distance_color(Z_cam)

        x1, y1, x2, y2 = obj['bbox_xyxy']

        cv2.rectangle(display_img, (x1, y1), (x2, y2), color, 2)
        cv2.putText(display_img,
                    f"{obj['class']} Z={Z_cam:.2f}",
                    (x1, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5, color, 2)

    # ======= 화면 출력 (고정 1280x720) =======
    resized_output = cv2.resize(display_img, (WINDOW_W, WINDOW_H))
    cv2.imshow("YOLO + Depth 3D", cv2.cvtColor(resized_output, cv2.COLOR_RGB2BGR))
    # ➕ 깊이맵 출력
    depth_resized = cv2.resize(depth_color, (WINDOW_W, WINDOW_H))
    cv2.imshow("Depth Map", depth_resized)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

 

다음 글에서는 위 통합 코드에 GUI를 추가하여 여러 기능을 커스터마이징 할 수 있도록 한 코드에 대해서 소개하겠습니다.