import tensorflow as tf
from keras import layers, models
import numpy as np
from keras.preprocessing import image
img_height = 128
img_width = 128
img_channels = 3
shape = (img_height, img_width, img_channels)
classes = 2
def create_model(input_shape, num_classes):
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=input_shape))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Flatten())
model.add(layers.Dense(128, activation='relu'))
model.add(layers.Dense(num_classes, activation='softmax'))
return model
model = create_model(shape, classes)
model.compile(optimizer=‘adam’,
loss=‘sparse_categorical_crossentropy’,
metrics=[‘accuracy’])
model.summary()
def parse_tfrecord_fn(example):
feature_description = {
‘image/encoded’: tf.io.FixedLenFeature(, tf.string),
‘image/filename’: tf.io.FixedLenFeature(, tf.string),
‘image/format’: tf.io.FixedLenFeature(, tf.string),
‘image/height’: tf.io.FixedLenFeature(, tf.int64),
‘image/width’: tf.io.FixedLenFeature(, tf.int64),
‘image/object/bbox/xmin’: tf.io.FixedLenFeature(, tf.float32),
‘image/object/bbox/xmax’: tf.io.FixedLenFeature(, tf.float32),
‘image/object/bbox/ymin’: tf.io.FixedLenFeature(, tf.float32),
‘image/object/bbox/ymax’: tf.io.FixedLenFeature(, tf.float32),
‘image/object/class/label’: tf.io.FixedLenFeature(, tf.int64),
‘image/object/class/text’: tf.io.FixedLenFeature(, tf.string),
}
example = tf.io.parse_single_example(example, feature_description)
filename = example['image/filename']
label = example['image/object/class/label']
print(f"Processing {filename}, label: {label}")
image = tf.image.decode_jpeg(example['image/encoded'], channels=3)
image = tf.image.resize(image, (img_height, img_width))
image = tf.image.convert_image_dtype(image, tf.float32)
return image, label
def create_dataset(tfrecord_path, batch_size, buffer_size):
dataset = tf.data.TFRecordDataset(tfrecord_path)
dataset = dataset.map(parse_tfrecord_fn)
dataset = dataset.shuffle(buffer_size)
dataset = dataset.batch(batch_size)
return dataset
train_dataset = create_dataset(‘default.tfrecord’, batch_size=12, buffer_size=1000)
test_dataset = create_dataset(‘default.tfrecord’, batch_size=12, buffer_size=1000)
model.fit(train_dataset, epochs=100, validation_data=test_dataset)
model.save(“model.keras”)
def load_label_map(label_map_path):
label_map = {}
with open(label_map_path, ‘r’) as file:
lines = file.readlines()
for line in lines:
if ‘id’ in line:
id_line = line.strip().split(‘:’)
label_id = int(id_line[1].strip())
elif ‘name’ in line:
name_line = line.strip().split(‘:’)
label_name = name_line[1].strip().replace(“'”, “”)
label_map[label_id] = label_name
return label_map
label_map_path = ‘label_map.pbtxt’
label_map = load_label_map(label_map_path)
img_path = ‘bmw.png’
img = image.load_img(img_path, target_size=(img_height, img_width))
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array /= 255.0
predictions = model.predict(img_array)
for i, prob in enumerate(predictions[0]):
label_id = i + 1
class_name = label_map.get(label_id, f’Unknown Class {label_id}')
print(f"{class_name}: {prob:.4f}")
When i have “classes = 2” i get accuracy, loss, etc equals to NaN. When i put a 3 there, it seems to be working, but of course i have a third class, that i dont really need.