Weights & Biases로 이미지 분류 최적화

import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import VGG16
from tensorflow.keras.layers import Dense, Flatten, Dropout
from tensorflow.keras.models import Model
from PIL import ImageFile
import wandb
from wandb.keras import WandbCallback

wandb.login()
wandb.init(project='wildfire_detection', config={
    "learning_rate": 0.001,
    "epochs": 5,
    "batch_size": 8,
    "target_size": (350, 350),
})
config = wandb.config

if tf.test.gpu_device_name():
    print('GPU device found:', tf.test.gpu_device_name())
else:
    print("No GPU found. Please ensure that GPU is enabled in the runtime settings.")

train_dir_wildfire = '/kaggle/input/wildfire-prediction-dataset/train'
valid_dir = '/kaggle/input/wildfire-prediction-dataset/valid'
target_size = config.target_size
batch_size = config.batch_size
seed = 42

ImageFile.LOAD_TRUNCATED_IMAGES = True

train_datagen_wildfire = ImageDataGenerator(
    rescale=1./255,
    rotation_range=40,      # Randomly rotate images in the range (degrees, 0 to 180)
    width_shift_range=0.2,  # Randomly horizontal shift images
    height_shift_range=0.2, # Randomly vertical shift images
    shear_range=0.2,        # Shear Intensity (Shear angle in counter-clockwise direction)
    zoom_range=0.2,         # Randomly zoom image
    horizontal_flip=True,   # Randomly flip images horizontally
    fill_mode='nearest'     # Strategy used for filling in newly created pixels
)

train_generator_wildfire = train_datagen_wildfire.flow_from_directory(
    directory=train_dir_wildfire,
    target_size=target_size,
    batch_size=batch_size,
    class_mode='binary',
    shuffle=True,
    seed=seed
)

validation_datagen = ImageDataGenerator(rescale=1./255)
validation_generator = validation_datagen.flow_from_directory(
    directory=valid_dir,
    target_size=target_size,
    batch_size=batch_size,
    class_mode='binary',
    shuffle=False
)

base_model = VGG16(weights='imagenet', include_top=False, input_shape=(target_size[0], target_size[1], 3))
x = base_model.output
x = Flatten()(x)
x = Dense(256, activation='relu')(x)
x = Dropout(0.5)(x)
predictions = Dense(1, activation='sigmoid')(x)
model = Model(inputs=base_model.input, outputs=predictions)
for layer in base_model.layers:
    layer.trainable = False
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

from tensorflow.keras.callbacks import Callback

class CustomWandbCallback(Callback):
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        wandb.log({
            'epoch': epoch,
            'loss': logs.get('loss'),
            'accuracy': logs.get('accuracy'),
            'val_loss': logs.get('val_loss'),
            'val_accuracy': logs.get('val_accuracy')
        })

model.fit(
    train_generator_wildfire,
    steps_per_epoch=steps_per_epoch,
    epochs=config.epochs,
    validation_data=validation_generator,
    validation_steps=validation_steps,
    callbacks=[CustomWandbCallback()]  # Use the custom callback here
)

test_loss, test_acc = model.evaluate(validation_generator, steps=validation_generator.samples // validation_generator.batch_size)
print('Validation accuracy:', test_acc)

wandb.log({'Validation Loss': test_loss, 'Validation Accuracy': test_acc})

model.save("wildfire_detection_model.h5")
wandb.save("wildfire_detection_model.h5")
wandb.Artifact("model", type="model", description="A model to detect wildfires from images")
wandb.log_artifact("wildfire_detection_model.h5", type="model", name="wildfire_detection_model")

test_dir_wildfire = '/kaggle/input/wildfire-prediction-dataset/test/wildfire'
test_dir_nowildfire = '/kaggle/input/wildfire-prediction-dataset/test/nowildfire'

from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing import image
import numpy as np
import os

def load_and_preprocess_images(dir_path, target_size):
    datagen = ImageDataGenerator(rescale=1./255)
    images = []
    labels = []  # Assuming 1 for wildfire and 0 for nowildfire for binary classification
    file_paths = []

    for file_name in os.listdir(dir_path):
        file_path = os.path.join(dir_path, file_name)
        img = image.load_img(file_path, target_size=target_size)
        img_array = image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_preprocessed = datagen.flow(img_array, batch_size=1)[0]

        images.append(img_preprocessed)
        file_paths.append(file_path)
    
    return np.vstack(images), file_paths

target_size = (350, 350)  # The target size used during training

images_wildfire, files_wildfire = load_and_preprocess_images(test_dir_wildfire, target_size)
images_nowildfire, files_nowildfire = load_and_preprocess_images(test_dir_nowildfire, target_size)

predictions_wildfire = model.predict(images_wildfire)

for i, pred in enumerate(predictions_wildfire):
    print(f"{files_wildfire[i]} - {'Wildfire' if pred >= 0.5 else 'No Wildfire'}: {pred[0]}")

predictions_nowildfire = model.predict(images_nowildfire)

for i, pred in enumerate(predictions_nowildfire):
    print(f"{files_nowildfire[i]} - {'Wildfire' if pred >= 0.5 else 'No Wildfire'}: {pred[0]}")

binary_predictions_wildfire = [1 if pred >= 0.5 else 0 for pred in predictions_wildfire]
binary_predictions_nowildfire = [1 if pred >= 0.5 else 0 for pred in predictions_nowildfire]

true_labels_wildfire = [1] * len(files_wildfire)  
true_labels_nowildfire = [0] * len(files_nowildfire)

def log_predictions_wandb(image_paths, predictions, true_labels, title, num_samples=10):
    columns = ["Image", "Predicted Label", "True Label"]
    wandb_table = wandb.Table(columns=columns)
    
    for i in range(min(num_samples, len(image_paths))):
        img = wandb.Image(image.load_img(image_paths[i], target_size=target_size))
        pred_label = "Wildfire" if predictions[i] >= 0.5 else "No Wildfire"
        true_label = "Wildfire" if true_labels[i] == 1 else "No Wildfire"
        wandb_table.add_data(img, pred_label, true_label)
        
    wandb.log({title: wandb_table})

og_predictions_wandb(files_wildfire, binary_predictions_wildfire, true_labels_wildfire, "Wildfire Predictions")
log_predictions_wandb(files_nowildfire, binary_predictions_nowildfire, true_labels_nowildfire, "Nowildfire Predictions")

wandb.finish()