본문 바로가기

AI

[객체검출] Object Detection #1 - 이미지 데이터 확인

 

  1. 라벨 데이터 읽기: 라벨 데이터의 경로를 찾아 불러와서 확인
  2. 이미지 데이터 읽기: 라벨 데이터와 같은 이름의 이미지 경로를 찾아 불러와서 확인
  3. 이미지 데이터 속성 파악: 이미지 데이터를 읽어 가로(width), 세로(height) 확인
  4. 라벨 데이터 속성 파악: 주어진 라벨 데이터는 바운딩 박스의 좌상단 정규화 좌표와 우하단 좌표의 정규화 좌표로 되어 있음을 확인
  5. 바운딩 박스 draw: 라벨 데이터 내 바운딩 박스의 좌표를 읽어 이미지 내 draw

 

import glob, os
from PIL import Image, ImageDraw, ImageFont

# 검출할 클래스 정보 리스트
class_idx = ['person', 'car', 'taxi', 'truck', 'bus', 'tricycle', 'motorcycle']

# 데이터 경로 지정
label_folder = f'/labels'
image_folder = f'/images'
result_folder = f'/results'

# 라벨(Label) 리스트 생성 및 갯수 출력
label_list = sorted(glob.glob(os.path.join(label_folder, "*.txt")))
print(len(label_list))

# 이미지 리스트 생성 및 갯수 출력
image_list = sorted(glob.glob(os.path.join(image_folder, "*.jpg")))
print(len(image_list))

# 라벨 데이터 속성 확인
# 라벨 데이터와 각각 매칭되는 이미지 데이터를 함께 불러와 BoundingBox 정화 확인 후 시각화 후 저장
# Bounding Box 표시된 이미지 결과 생성

count = 0 # 현재 몇개의 데이터를 처리하는 지 count하는 변수 선언

for label_file in label_list:

    name = label_file.split('/')[-1] # 각 라벨 데이터와 매칭되는 이미지 데이터의 이름을 확인하기 위해 현재 라벨데이터의 이름 문자열 parsing
    image_name = image_folder + '/' + name.replace('txt','jpg') 
    write_image_name = result_folder + '/' + name.replace('txt','jpg') # 바운딩 박스를 그린 후 결과를 이미지 형태로 저장할 때 원래 데이터 파일명과 동일하게 저장

    img = Image.open(image_name).convert('RGB')
    w, h = img.size # 이미지의 가로와 세로 길이 추출

    with open(label_file, 'r') as f: # 라벨 데이터 읽기 (with ~ as 구문 사용시 f.close() 부분은 생략 가능함)
        lines = f.read().splitlines()

    for line in lines:
        cls_idx, xmin, ymin, xmax, ymax = line.split(' ') # 각 라벨 파일은 클래스 정보, 바운딩박스의 좌상단 x,y, 우하단 x,y 의 정규화된 좌표가 각 줄마다 적혀있으며, 각 정보는 띄어쓰기로 구분되어 있음

        xmin, ymin, xmax, ymax = map(float, [xmin, ymin, xmax, ymax]) # 클래스 정보를 제외한 모든 좌표의 값은 float 형태로 변환
        
        # 라벨 데이터의 바운딩박스 정보들은 0과 1사이의 값으로 정규화되어 있기 때문에 위에서 획득한 이미지의 가로,세로 정보를 활용하여 실제 좌표값으로 변환
        xmin *= w
        ymin *= h
        xmax *= w
        ymax *= h
        cls_name = class_idx[int(cls_idx)] # 클래스 정보 (이름)는 0부터 시작하는 클래스 인덱스 값으로 획득 가능

        # 이미지 위에 rectangle 함수를 이용하여 (좌상단 좌표, 우하단 좌표)로 그린 후, 클래스 정보를 함께 표기
        draw = ImageDraw.Draw(img)
        draw.rectangle((xmin,ymin, xmax, ymax), outline=(0,255,0), width=3) 
        draw.text((xmin,ymin-10), cls_name, (255,0,255))

    # 이미지위에 모든 객체의 바운딩 박스를 그려 위에서 정의한 경로에 저장
    img.save(write_image_name)
    count += 1
    print("%d images / total %d images processed"%(count,len(image_list))) #이미지 처리 프로세스 현황 출력
    
    f.close() # optional (with ~ as 구문으로 파일을 열었기에 반드시 필요한 부분은 아님)