def main():
if not exists('results'):
os.makedirs('results')
file_names = []
with open(join(config.base_dir, 'val.txt')) as reader:
lines = reader.readlines()
for line in lines:
file_names.append(line.rstrip().split(' ')[0])
model = nn.build_model(training=False)
model.load_weights(f"weights/model_{config.version}.h5", True)
for file_name in tqdm.tqdm(file_names):
image = cv2.imread(join(config.base_dir, config.image_dir, file_name + '.jpg'))
image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_np, scale, dw, dh = util.resize(image_np)
image_np = image_np.astype(np.float32) / 255.0
boxes, scores, _ = model.predict(image_np[np.newaxis, ...])
boxes, scores = np.squeeze(boxes, 0), np.squeeze(scores, 0)
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / scale
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / scale
image = draw_bbox(image, boxes, scores)
cv2.imwrite(f'results/{file_name}.jpg', image)
def main():
score_threshold = 0.00001
if not exists(join('results', f'D{config.phi}')):
os.makedirs(join('results', f'D{config.phi}'))
train_model, inference_model = build_model(phi=config.phi,
num_classes=len(config.classes),
score_threshold=score_threshold)
train_model.load_weights(config.weight_path)
train_model.save_weights(config.weight_path)
inference_model.load_weights(config.weight_path, by_name=True)
f_names = []
with open(join(config.data_dir, 'val.txt')) as reader:
for line in reader.readlines():
f_names.append(line.rstrip().split(' ')[0])
result_dict = {}
for f_name in tqdm.tqdm(f_names):
image_path = join(config.data_dir, config.image_dir, f_name + '.jpg')
label_path = join(config.data_dir, config.label_dir, f_name + '.xml')
image = cv2.imread(image_path)
src_image = image.copy()
image = image[:, :, ::-1]
h, w = image.shape[:2]
image, scale = preprocess_image(image, image_size=config.image_size)
boxes, scores, labels = inference_model.predict_on_batch(
[np.expand_dims(image, axis=0)])
boxes, scores, labels = np.squeeze(boxes), np.squeeze(
scores), np.squeeze(labels)
boxes = postprocess_boxes(boxes=boxes, scale=scale, height=h, width=w)
indices = np.where(scores[:] > score_threshold)[0]
boxes = boxes[indices]
draw_boxes(src_image, boxes)
pred_boxes_np = []
for pred_box in boxes:
x_min, y_min, x_max, y_max = pred_box
pred_boxes_np.append([x_min, y_min, x_max, y_max])
true_boxes = []
for element in parse_fn(label_path).getroot().iter('object'):
box = find_node(element, 'bndbox')
x_min = find_node(box, 'xmin', 'bndbox.xmin', parse=float) - 1
y_min = find_node(box, 'ymin', 'bndbox.ymin', parse=float) - 1
x_max = find_node(box, 'xmax', 'bndbox.xmax', parse=float) - 1
y_max = find_node(box, 'ymax', 'bndbox.ymax', parse=float) - 1
true_boxes.append([x_min, y_min, x_max, y_max])
result = {
'detection_boxes': pred_boxes_np,
'groundtruth_boxes': true_boxes,
'confidence': scores
}
result_dict[f'{f_name}.jpg'] = result
cv2.imwrite(join('results', f'D{config.phi}', basename(image_path)),
src_image[:, :, ::-1])
# cv2.namedWindow('image', cv2.WINDOW_NORMAL)
# cv2.imshow('image', src_image)
# cv2.waitKey(0)
with open(join('results', f'D{config.phi}', 'd4.pickle'), 'wb') as writer:
pickle.dump(result_dict, writer)
def test_fn():
if not os.path.exists('results'):
os.makedirs('results')
file_names = []
with open(os.path.join(config.data_dir, 'test.txt')) as f:
for file_name in f.readlines():
file_names.append(file_name.rstrip())
model = nn.build_model(training=False)
model.load_weights(f"weights/model_{config.version}.h5", True)
for file_name in tqdm.tqdm(file_names):
image = cv2.imread(
os.path.join(config.data_dir, config.image_dir,
file_name + '.jpg'))
image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_np, scale, dw, dh = util.resize(image_np)
image_np = image_np.astype(numpy.float32) / 255.0
boxes, scores, labels = model.predict(image_np[numpy.newaxis, ...])
boxes, scores, labels = numpy.squeeze(boxes, 0), numpy.squeeze(
scores, 0), numpy.squeeze(labels, 0)
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / scale
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / scale
image = draw_bbox(image, boxes)
cv2.imwrite(f'results/{file_name}.jpg', image)
def main():
nb_classes = len(config.classes)
input_layer = tf.keras.layers.Input([config.image_size, config.image_size, 3])
feature_maps = nn.build_model(input_layer, nb_classes)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = nn.decode(fm, i, nb_classes)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.load_weights("weights/model.h5")
image = cv2.imread('../Dataset/VOC2012/IMAGES/2007_000039.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_size = image.shape[:2]
image_data = util.resize(image)
pred_bbox = model.predict(image_data[np.newaxis, ...].astype(np.float32) / 255.0)
pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
pred_bbox = tf.concat(pred_bbox, axis=0)
bbox = postprocess_boxes(pred_bbox, image_size, config.image_size, 0.3)
bbox = nms(bbox)
image = draw_bbox(image, bbox, nb_classes)
cv2.imwrite('result.png', image)
def main():
model = nn.build_model(training=False)
model.load_weights("weights/model2.h5", True)
f_names = []
with open(join(config.data_dir, 'val.txt')) as f:
for f_name in f.readlines()[:100]:
f_names.append(f_name.rstrip().split(' ')[0])
for f_name in f_names[:1]:
image = cv2.imread(
join(config.data_dir, config.image_dir, f'{f_name}.jpg'))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_data, scale, dw, dh = util.resize(image)
image_data = image_data.astype(np.float32) / 255.0
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
image_data -= mean
image_data /= std
boxes, score, label = model.predict(image_data[np.newaxis, ...])
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / scale
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / scale
image = draw_bbox(image, boxes)
cv2.imwrite(f'results/{f_name}.png',
cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
def main():
_, palette = util.get_label_info(
os.path.join(config.data_dir, "class_dict.csv"))
model = nn.build_model(classes=len(palette))
model(tf.zeros((1, config.height, config.width, 3)))
file_names = [
file_name[:-4] for file_name in os.listdir(
os.path.join(config.data_dir, config.image_dir))
]
for file_name in tqdm.tqdm(file_names):
image = util.load_image(file_name)
label = util.load_label(file_name)
image, label = util.random_crop(image, label)
image = np.expand_dims(image, 0).astype('float32')
output = model.predict(image / 255.0)
output = np.array(output[0, :, :, :])
output = np.argmax(output, axis=-1)
output = util.colour_code_segmentation(output, palette)
output = np.uint8(output)
util.save_images([output, label],
os.path.join('results', f'{file_name}.jpg'),
titles=['Pred', 'Label'])
def main():
if not exists('results'):
os.makedirs('results')
f_names = []
with open(join(config.base_dir, 'val.txt')) as reader:
lines = reader.readlines()
for line in lines:
f_names.append(line.rstrip().split(' ')[0])
result_dict = {}
model = nn.build_model(training=False)
model.load_weights("weights/model245.h5", True)
for f_name in tqdm.tqdm(f_names):
image_path = join(config.base_dir, config.image_dir, f_name + '.jpg')
label_path = join(config.base_dir, config.label_dir, f_name + '.xml')
image = cv2.imread(image_path)
image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_np, scale, dw, dh = util.resize(image_np)
image_np = image_np.astype(np.float32) / 255.0
boxes, scores, _ = model.predict(image_np[np.newaxis, ...])
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / scale
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / scale
true_boxes = []
for element in parse_fn(label_path).getroot().iter('object'):
true_boxes.append(parse_annotation(element)[2])
result = {
'pred_boxes': boxes,
'true_boxes': true_boxes,
'confidence': scores
}
result_dict[f'{f_name}.jpg'] = result
image = draw_bbox(image, boxes, scores)
cv2.imwrite(f'results/{f_name}.png', image)
with open(join('results/yolo.pickle'), 'wb') as writer:
pickle.dump(result_dict, writer)
box_classes = backend.argmax(box_scores, -1)
box_class_scores = backend.max(box_scores, -1)
prediction_mask = box_class_scores >= score_threshold
_boxes = tf.boolean_mask(_boxes, prediction_mask)
scores = tf.boolean_mask(box_class_scores, prediction_mask)
_classes = tf.boolean_mask(box_classes, prediction_mask)
_boxes = _boxes * config.image_size
selected_idx = tf.image.non_max_suppression(_boxes, scores, config.max_boxes, iou_threshold)
return backend.gather(_boxes, selected_idx), backend.gather(_classes, selected_idx)
if __name__ == '__main__':
inputs = tf.keras.layers.Input((config.image_size, config.image_size, 3))
outputs = nn.build_model(inputs, len(config.classes), False)
model = tf.keras.models.Model(inputs, outputs)
model.load_weights(join('weights', 'model15.h5'))
path = '../Dataset/VOC2012/IMAGES/2007_000027.jpg'
image = cv2.imread(path)
image = image[:, :, ::-1]
image = util.resize(image)
input_image = image.astype('float32') / 127.5 - 1.0
input_image = np.expand_dims(input_image, 0)
boxes, classes = get_box(model.predict_on_batch(input_image), 0.45, 0.4)
for i, box in enumerate(boxes):
x1 = int(box[0])
y1 = int(box[1])
x2 = int(box[2])
y2 = int(box[3])
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 1)
f_names = []
with open(join(config.base_dir, 'train.txt')) as reader:
for line in reader.readlines():
f_names.append(line.rstrip().split(' ')[0])
num_replicas = strategy.num_replicas_in_sync
steps = len(f_names) // config.batch_size
lr = nn.CosineLrSchedule(steps)
nb_classes = len(config.classes)
dataset = data_loader.input_fn(f_names)
dataset = strategy.experimental_distribute_dataset(dataset)
with strategy.scope():
model = nn.build_model()
optimizer = tf.keras.optimizers.Adam(lr, beta_1=0.935, decay=0.0005)
with strategy.scope():
loss_object = nn.compute_loss
def compute_loss(y_true, y_pred):
total_loss = loss_object(y_pred, y_true)
return tf.reduce_sum(total_loss) * 1.0 / num_replicas
def train_step(image, y_true):
with tf.GradientTape() as tape:
y_pred = model(image, training=True)
loss = compute_loss(y_true, y_pred)
strategy = tf.distribute.MirroredStrategy()
_, palette = util.get_label_info(
os.path.join(config.data_dir, 'class_dict.csv'))
file_names = [
file_name[:-4] for file_name in os.listdir(
os.path.join(config.data_dir, config.image_dir))
]
dataset = input_fn(file_names, palette)
dataset = strategy.experimental_distribute_dataset(dataset)
weights = util.get_class_weights(file_names)
with strategy.scope():
optimizer = tf.keras.optimizers.RMSprop(0.0001, decay=0.995)
model = nn.build_model((config.height, config.width, 3), len(palette))
model(tf.zeros((1, config.height, config.width, 3)))
with strategy.scope():
loss_fn = nn.segmentation_loss(weights)
def compute_loss(y_true, y_pred):
return tf.reduce_sum(loss_fn(y_true, y_pred)) * 1. / config.batch_size
with strategy.scope():
def train_step(image, y_true):
with tf.GradientTape() as tape:
y_pred = model(image)
loss = compute_loss(y_true, y_pred)
strategy = tf.distribute.MirroredStrategy()
nb_gpu = strategy.num_replicas_in_sync
global_batch = nb_gpu * config.batch_size
nb_classes = len(config.classes)
dataset = data_loader.TFRecordLoader(global_batch, config.epochs,
nb_classes).load_data(tf_paths)
dataset = strategy.experimental_distribute_dataset(dataset)
with strategy.scope():
optimizer = tf.keras.optimizers.RMSprop(1e-4)
input_tensor = tf.keras.layers.Input(
[config.image_size, config.image_size, 3])
outputs = nn.build_model(input_tensor, nb_classes)
output_tensors = []
for i, output in enumerate(outputs):
pred_tensor = nn.decode(output, i, nb_classes)
output_tensors.append(output)
output_tensors.append(pred_tensor)
model = tf.keras.Model(input_tensor, output_tensors)
print(f'[INFO] {len(tf_paths)} train data')
with strategy.scope():
loss_object = nn.compute_loss
def compute_loss(y_true, y_pred):
iou_loss = conf_loss = prob_loss = 0
def train():
strategy = tf.distribute.MirroredStrategy()
file_names = []
with open(os.path.join(config.data_dir, 'train.txt')) as f:
for file_name in f.readlines():
image_path = os.path.join(config.data_dir, config.image_dir,
file_name.rstrip() + '.jpg')
label_path = os.path.join(config.data_dir, config.label_dir,
file_name.rstrip() + '.xml')
if os.path.exists(image_path) and os.path.exists(label_path):
if os.path.exists(os.path.join(config.data_dir, 'TF')):
file_names.append(
os.path.join(config.data_dir, 'TF',
file_name.rstrip() + '.tf'))
else:
file_names.append(file_name.rstrip())
steps = len(file_names) // config.batch_size
if os.path.exists(os.path.join(config.data_dir, 'TF')):
dataset = DataLoader().input_fn(file_names)
else:
dataset = input_fn(file_names)
dataset = strategy.experimental_distribute_dataset(dataset)
with strategy.scope():
model = nn.build_model()
model.summary()
optimizer = tf.keras.optimizers.Adam(nn.CosineLR(steps), 0.937)
with strategy.scope():
loss_object = nn.ComputeLoss()
def compute_loss(y_true, y_pred):
total_loss = loss_object(y_pred, y_true)
return tf.reduce_sum(total_loss) / config.batch_size
with strategy.scope():
def train_step(image, y_true):
with tf.GradientTape() as tape:
y_pred = model(image, training=True)
loss = compute_loss(y_true, y_pred)
variables = model.trainable_variables
gradients = tape.gradient(loss, variables)
optimizer.apply_gradients(zip(gradients, variables))
return loss
with strategy.scope():
@tf.function
def distributed_train_step(image, y_true):
per_replica_losses = strategy.run(train_step, args=(image, y_true))
return strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_losses,
axis=None)
def train_fn():
if not os.path.exists('weights'):
os.makedirs('weights')
pb = tf.keras.utils.Progbar(steps, stateful_metrics=['loss'])
print(f'[INFO] {len(file_names)} data points')
for step, inputs in enumerate(dataset):
if step % steps == 0:
print(f'Epoch {step // steps + 1}/{config.num_epochs}')
pb = tf.keras.utils.Progbar(steps, stateful_metrics=['loss'])
step += 1
image, y_true_1, y_true_2, y_true_3 = inputs
y_true = (y_true_1, y_true_2, y_true_3)
loss = distributed_train_step(image, y_true)
pb.add(1, [('loss', loss)])
if step % steps == 0:
model.save_weights(
os.path.join("weights", f"model_{config.version}.h5"))
if step // steps == config.num_epochs:
sys.exit("--- Stop Training ---")
train_fn()
