Computer Vision - LAB 05. Image Classification (OpenCV, Keras)

 

간단한 이미지 분류 프로젝트를 진행해보자. 이번 글에서는 OpenCV로 이미지를 처리하고, Keras로 모델을 학습시켜 심슨 가족의 Homer와 Bart 이미지를 구분하는 신경망을 만들어 보겠다.

---

1. 이미지 데이터 전처리

먼저, 이미지를 읽고 학습할 수 있는 벡터 형태로 변환하는 전처리 과정을 거친다. 디렉토리에서 파일을 불러와 벡터로 변환하고, 각 벡터에 해당하는 클래스 레이블을 설정한다.

import cv2
import numpy as np
import os
import keras
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

##############################################################################
# 이미지 파일 경로 설정 및 파일 목록 생성
# - 이미지 파일이 저장된 디렉토리를 지정하고, 해당 디렉토리 내 모든 파일의 경로를 리스트로 생성한다.
directory = r"D:\Project\COMPUTER_VISION\PYTHON\CNN\homer_bart_1"
files = [os.path.join(directory, f) for f in sorted(os.listdir(directory))]

#############################################################################
# 이미지 파일을 벡터로 변환하고 클래스 레이블 추출
# - 이미지 파일을 읽어 벡터 형태로 변환하고, 파일 이름에 따라 클래스(0 또는 1)를 지정한다.
height, width = 128, 128
images = []
classes = []

for image_path in files:

    # 이미지 읽기 및 오류 처리
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: Could not load image at {image_path}")
        continue

    # 이미지 전처리: 크기 조정 및 그레이스케일 변환
    image = cv2.resize(image, (width, height))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = image.ravel()  # 2D 배열(128x128)을 1D 배열(16384)로 평탄화
    images.append(image)


    # 클래스 레이블 설정
    image_name = os.path.basename(os.path.normpath(image_path))  # 파일 이름 추출
    if image_name.startswith("b"):
        class_name = 0
    else:
        class_name = 1
    classes.append(class_name)

---

2. 이미지 데이터 정규화, 분할

신경망 학습을 위해 데이터를 정규화하고, 학습/테스트 셋으로 나눈다.

from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

scaler = MinMaxScaler()  # MinMaxScaler를 사용해 데이터를 0~1 범위로 정규화

X = np.asarray(images)  # 이미지 리스트를 NumPy 배열로 변환
Y = np.asarray(classes)  # 클래스 리스트를 NumPy 배열로 변환

X = scaler.fit_transform(X)  # 0~1 사이로 정규화

# 데이터셋을 학습용(80%)과 테스트용(20%)으로 분할
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=1)

---

3. 신경망 모델 설계

두 이미지 중 하나를 분류하는 이진 분류이기 때문에, Fully Connected Layer를 갖는 MLP(다층 퍼셉트론)를 만든다.

model = keras.models.Sequential()  # 순차적 모델 생성

# 입력층: 16384개의 픽셀 입력 (128x128), 출력 8193개 노드 [(input + output) / 2], ReLU 활성화 함수
model.add(keras.layers.Dense(input_shape=(16384,), units=8193, activation="relu"))
# 은닉층: 8193개 노드, ReLU 활성화 함수
model.add(keras.layers.Dense(units=8193, activation="relu"))
# 출력층: 이진 분류를 위한 1개 노드, Sigmoid 활성화 함수
model.add(keras.layers.Dense(units=1, activation="sigmoid"))

print(model.summary())  # 모델 구조 출력

 

---

4. 모델 학습 및 저장

모델을 50 에포크 동안 학습시켜 보고, 나중에 재사용할 수 있도록 모델을 저장한다. Scaler 또한 저장하여 Inference 때 재사용할 수 있게 만든다.

model.fit(X_train, Y_train, epochs=50)

######################################################
# 1. 모델 전체를 저장 (구조, 가중치, 옵티마이저 포함)
model.save("./Network/homer_bart_model.keras")

# 2. 가중치만 따로 저장
model.save_weights("./Network/homer_bart.weights.h5")

# 3. Scaler 저장
with open("./Network/scaler.pkl", "wb") as f:
    pickle.dump(scaler, f)
print("Scaler saved to ./Network/scaler.pkl")

 

Epoch를 진행할 수록 정확도가 올라가는 모습을 확인할 수 있다.

---

5. Inference

1 ~ 4 과정을 통해 학습된 Model을 Load하고, predict() 함수를 통해 이미지를 잘 분류할 수 있는지 시험해본다.

import cv2
import numpy as np
import keras
import pickle

# predict 함수 : image를 전처리하고, 그에 맞는 결과를 출력 (0: Bart, 1: Homer)
def predict(image_path):
    image = cv2.imread(image_path)
    height, width = 128, 128

    if image is None:
        print(f"Error: Could not load image at {image_path}")
        return None

    image = cv2.resize(image, (width, height))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = image.ravel()

    image = np.asarray([image])
    image = scaler.transform(image)

    prediction = model.predict(image)
    result = 1 if prediction[0][0] >= 0.5 else 0

    return result


# 모델 로드
model = keras.models.load_model("./Network/homer_bart_model.keras")

# 학습 때 저장한 scaler 로드
with open("./Network/scaler.pkl", "rb") as f:
    scaler = pickle.load(f)


for i in range(1,8):
    test_image_path = r"D:\Project\COMPUTER_VISION\PYTHON\CNN\homer_bart_1\bart"
    result = predict(f"{test_image_path}{i}.bmp")
    print(f"Prediction for {test_image_path}{i}: {result}")


for i in range(27,33):
    test_image_path = r"D:\Project\COMPUTER_VISION\PYTHON\CNN\homer_bart_1\homer"
    result = predict(f"{test_image_path}{i}.bmp")
    print(f"Prediction for {test_image_path}{i}: {result}")


"""
전체 모델이 아닌 가중치만 Load해서 사용할 경우, 모델 구조를 다시 정의해줘야 함!!


# 모델 구조 정의 (저장 전과 동일해야 함)
model = keras.models.Sequential()
model.add(keras.layers.Dense(input_shape=(16384,), units=8193, activation="relu"))
model.add(keras.layers.Dense(units=8193, activation="relu"))
model.add(keras.layers.Dense(units=1, activation="sigmoid"))

# 가중치 로드
model.load_weights("./Network/homer_bart_weights.weights.h5")

# 컴파일 (필요 시)
model.compile(optimizer="Adam", loss="binary_crossentropy", metrics=["accuracy"])

# 모델 확인
model.summary()
"""

---

* Code for Training

import cv2
import numpy as np
import os
import keras
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

##############################################################################
# 이미지 파일 경로 설정 및 파일 목록 생성
# - 이미지 파일이 저장된 디렉토리를 지정하고, 해당 디렉토리 내 모든 파일의 경로를 리스트로 생성한다.
directory = r"D:\Project\COMPUTER_VISION\PYTHON\CNN\homer_bart_1"
files = [os.path.join(directory, f) for f in sorted(os.listdir(directory))]

#############################################################################
# 이미지 파일을 벡터로 변환하고 클래스 레이블 추출
# - 이미지 파일을 읽어 벡터 형태로 변환하고, 파일 이름에 따라 클래스(0 또는 1)를 지정한다.
height, width = 128, 128
images = []
classes = []

for image_path in files:

    # 이미지 읽기 및 오류 처리
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: Could not load image at {image_path}")
        continue

    # 이미지 전처리: 크기 조정 및 그레이스케일 변환
    image = cv2.resize(image, (width, height))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = image.ravel()  # 2D 배열(128x128)을 1D 배열(16384)로 평탄화
    images.append(image)


    # 클래스 레이블 설정
    image_name = os.path.basename(os.path.normpath(image_path))  # 파일 이름 추출
    if image_name.startswith("b"):
        class_name = 0
    else:
        class_name = 1
    classes.append(class_name)

############################################################################
# 데이터 정규화 및 학습/테스트 데이터 분할
# - 이미지 데이터를 정규화하고, 학습 데이터와 테스트 데이터로 나눈다.
scaler = MinMaxScaler()  # MinMaxScaler를 사용해 데이터를 0~1 범위로 정규화

X = np.asarray(images)  # 이미지 리스트를 NumPy 배열로 변환
Y = np.asarray(classes)  # 클래스 리스트를 NumPy 배열로 변환

X = scaler.fit_transform(X)  # 이미지 데이터를 정규화 (0~1 사이 값으로 스케일링)

# 데이터셋을 학습용(80%)과 테스트용(20%)으로 분할
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=1)

##############################################################################
# 신경망 모델 정의 및 학습
# - Keras를 사용해 신경망 모델을 정의하고, 컴파일 후 학습을 수행한다.
model = keras.models.Sequential()  # 순차적 모델 생성

# 입력층: 16384개의 픽셀 입력 (128x128), 출력 8193개 노드 [(input + output) / 2], ReLU 활성화 함수
model.add(keras.layers.Dense(input_shape=(16384,), units=8193, activation="relu"))
# 은닉층: 8193개 노드, ReLU 활성화 함수
model.add(keras.layers.Dense(units=8193, activation="relu"))
# 출력층: 이진 분류를 위한 1개 노드, Sigmoid 활성화 함수
model.add(keras.layers.Dense(units=1, activation="sigmoid"))

print(model.summary())  # 모델 구조 출력

# 모델 컴파일: Adam 옵티마이저, 이진 분류 손실 함수, 정확도 메트릭 사용
model.compile(optimizer="Adam", loss="binary_crossentropy", metrics=["accuracy"])

# 모델 학습: 학습 데이터로 50 에포크 동안 학습
model.fit(X_train, Y_train, epochs=50)

######################################################
# 1. 모델 전체를 저장 (구조, 가중치, 옵티마이저 포함)
model.save("./Network/homer_bart_model.keras")

# 2. 가중치만 따로 저장
model.save_weights("./Network/homer_bart.weights.h5")

---

* Code for Inference

import cv2
import numpy as np
import keras
import pickle

# predict 함수 : image를 전처리하고, 그에 맞는 결과를 출력 (0: Bart, 1: Homer)
def predict(image_path):
    image = cv2.imread(image_path)
    height, width = 128, 128

    if image is None:
        print(f"Error: Could not load image at {image_path}")
        return None

    image = cv2.resize(image, (width, height))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = image.ravel()

    image = np.asarray([image])
    image = scaler.transform(image)

    prediction = model.predict(image)
    result = 1 if prediction[0][0] >= 0.5 else 0

    return result


# 모델 로드
model = keras.models.load_model("./Network/homer_bart_model.keras")

# 학습 때 저장한 scaler 로드
with open("./Network/scaler.pkl", "rb") as f:
    scaler = pickle.load(f)


for i in range(1,8):
    test_image_path = r"D:\Project\COMPUTER_VISION\PYTHON\CNN\homer_bart_1\bart"
    result = predict(f"{test_image_path}{i}.bmp")
    print(f"Prediction for {test_image_path}{i}: {result}")


for i in range(27,33):
    test_image_path = r"D:\Project\COMPUTER_VISION\PYTHON\CNN\homer_bart_1\homer"
    result = predict(f"{test_image_path}{i}.bmp")
    print(f"Prediction for {test_image_path}{i}: {result}")