File size: 3,048 Bytes
9333188 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import os
import shutil
import random
import numpy as np
import json
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
SOURCE_DIR = r"E:\Dog_Breed_Classification_model\dataset"
TARGET_DIR = r"E:\Dog_Breed_Classification_model\dataset_split"
IMG_SIZE = (224, 224)
BATCH_SIZE = 16
EPOCHS = 20
SEED = 42
def split_dataset():
random.seed(SEED)
for split in ['train', 'val']:
os.makedirs(os.path.join(TARGET_DIR, split), exist_ok=True)
for breed in os.listdir(SOURCE_DIR):
class_path = os.path.join(SOURCE_DIR, breed)
if not os.path.isdir(class_path):
continue
images = [img for img in os.listdir(class_path) if img.lower().endswith(('.jpg', '.jpeg', '.png'))]
train_imgs, val_imgs = train_test_split(images, test_size=0.2, random_state=SEED)
for split, img_list in zip(['train', 'val'], [train_imgs, val_imgs]):
split_dir = os.path.join(TARGET_DIR, split, breed)
os.makedirs(split_dir, exist_ok=True)
for img in img_list:
shutil.copy2(os.path.join(class_path, img), os.path.join(split_dir, img))
def create_generators():
train_gen = ImageDataGenerator(rescale=1./255, rotation_range=15, width_shift_range=0.1,
height_shift_range=0.1, horizontal_flip=True)
val_gen = ImageDataGenerator(rescale=1./255)
train_data = train_gen.flow_from_directory(os.path.join(TARGET_DIR, 'train'),
target_size=IMG_SIZE, batch_size=BATCH_SIZE,
class_mode='categorical')
val_data = val_gen.flow_from_directory(os.path.join(TARGET_DIR, 'val'),
target_size=IMG_SIZE, batch_size=BATCH_SIZE,
class_mode='categorical')
return train_data, val_data
def build_model(num_classes):
base_model = tf.keras.applications.MobileNetV2(input_shape=(*IMG_SIZE, 3), include_top=False, weights='imagenet')
base_model.trainable = False
model = tf.keras.Sequential([
base_model,
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.3),
tf.keras.layers.Dense(num_classes, activation='softmax')
])
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
return model
def save_model(model, class_indices):
model.save("dog_breed_classifier.h5")
with open("class_indices.json", "w") as f:
json.dump(class_indices, f)
if __name__ == "__main__":
split_dataset()
train_data, val_data = create_generators()
model = build_model(num_classes=len(train_data.class_indices))
model.fit(train_data, validation_data=val_data, epochs=EPOCHS)
save_model(model, train_data.class_indices) |