tensorflow 이미지 분류 모델

2023.12.21 - [컴퓨터/TensorFlow] - tensorflow 이미지 데이터셋 만들기

tensorflow 이미지 데이터셋 만들기

---

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

def load_and_preprocess_image(image_path):
    # 이미지 파일을 불러오고 전처리하는 함수
    image = tf.io.read_file(image_path)  # 이미지 파일 읽기
    image = tf.image.decode_jpeg(image, channels=1)  # 이미지를 흑백으로 디코딩
    image = tf.image.resize(image, [28, 28])  # 이미지 크기를 28x28로 조절
    image /= 255.0  # 픽셀 값을 0~1 범위로 정규화
    return image

# npz 파일 로드
data = np.load('datasets.npz')  # 저장된 npz 파일에서 데이터 로드

# 데이터를 키에 맞춰 불러오기
train_images = data['train_images']  # 훈련 이미지
train_labels = data['train_labels']  # 훈련 라벨
test_images = data['test_images']    # 테스트 이미지
test_labels = data['test_labels']    # 테스트 라벨

# 데이터셋을 TensorFlow 데이터셋으로 변환
train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(buffer_size=1000).batch(32)  # 훈련 데이터셋
test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(32)  # 테스트 데이터셋

# 클래스 이름 정의
class_names = ['a', 'b']  # 클래스 a와 b

# 모델 정의
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),  # 입력 이미지를 1차원 배열로 변환
    tf.keras.layers.Dense(4, activation='relu'),    # 4개의 유닛을 가진 Dense 층
    tf.keras.layers.Dense(2)                        # 최종 출력을 위한 2개의 유닛을 가진 Dense 층
])

# 모델 컴파일
model.compile(optimizer='adam',  # 옵티마이저로 'adam' 사용
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),  # 손실 함수
              metrics=['accuracy'])  # 평가 지표로 'accuracy' 사용

# 모델 학습
model.fit(train_dataset, epochs=3)  # 모델을 훈련 데이터셋으로 3번 에포크 동안 학습

# 모델 평가
test_loss, test_acc = model.evaluate(test_dataset, verbose=2)  # 테스트 데이터셋으로 모델 평가
print(f"\nTest accuracy: {test_acc}")  # 테스트 정확도 출력

# 새 이미지 불러오기 및 전처리
image = load_and_preprocess_image('image.jpg')  # 새 이미지 불러오기 및 전처리

# 이미지를 배치 형태로 만듦
image_batch = tf.expand_dims(image, 0)  # 배치 차원 추가

# 모델을 사용하여 예측
predictions = model.predict(image_batch)  # 이미지에 대한 예측 수행

# 예측 결과 해석
predicted_label = class_names[np.argmax(predictions[0])]  # 가장 높은 확률을 가진 클래스 선택
print(f"이 이미지는 '{predicted_label}' 클래스로 예측됩니다.")  # 예측된 클래스 출력