def __init__(self, model_path, weights_path, gpu_num=1):
if not self.check_model(model_path):
self.check_weights(weights_path)
self.score = config.score
self.iou = config.iou
self.weights_path = weights_path
self.model_path = model_path
self.gpu_num = gpu_num
self.colors = utils_image.get_random_colors(
len(config.classes_names))
self.yolo4_model = YOLO(config.image_input_shape)()
self.convertor()
print('Converting finished !')
def load_model(session, m_type, m_name, logger):
# load the weights based on best loss
best_dir = "best_loss"
# check model dir
model_path = "models/" + m_name
path = os.path.join(model_path, best_dir)
if not os.path.exists(path):
raise FileNotFoundError
if m_type == "simple":
model = Simple(m_name, config, logger)
elif m_type == "YOLO":
model = YOLO(m_name, config, logger)
elif m_type == "GAP":
model = GAP(m_name, config, logger)
elif m_type == "NAS":
model = NASNET(m_name, config, logger)
elif m_type == "INC":
model = Inception(m_name, config, logger)
else:
raise ValueError
# load the best saved weights
ckpt = tf.train.get_checkpoint_state(path)
if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
logger.log('Reloading model parameters..')
model.restore(session, ckpt.model_checkpoint_path)
else:
raise ValueError('There is no best model with given model')
return model
def create_model(session, m_type, m_name, logger):
"""
create or load the last saved model
:param session: tf.session
:param m_type: model type
:param m_name: model name (equal to folder name)
:param logger: logger
:return: None
"""
if m_type == "simple":
model = Simple(m_name, config, logger)
elif m_type == "YOLO":
model = YOLO(m_name, config, logger)
elif m_type == 'GAP':
model = GAP(m_name, config, logger)
elif m_type == 'NAS':
model = NASNET(m_name, config, logger)
elif m_type == 'INC':
model = Inception(m_name, config, logger)
else:
raise ValueError
ckpt = tf.train.get_checkpoint_state(model.model_dir)
if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
logger.log('Reloading model parameters..')
model.restore(session, ckpt.model_checkpoint_path)
else:
logger.log('Created new model parameters..')
session.run(tf.global_variables_initializer())
return model
def __init__(self, args):
print("args = {")
for k in args:
print("\t{} = {}".format(k, args[k]))
print("}")
self.args = args.copy()
if self.args["dataset_path"] != "":
self.data_loader = VOCDataLoader(self.args["dataset_path"],
num_processes=4,
preload=self.args["preload"])
self.classes = read_classes(self.args["class_path"])
self.color_dict = get_color_dict(self.classes, self.args["color_path"])
if self.args["model_save_dir"] != "" and not os.path.exists(
self.args["model_save_dir"]):
os.makedirs(self.args["model_save_dir"])
if self.args["graph_save_dir"] != "" and not os.path.exists(
self.args["graph_save_dir"]):
os.makedirs(self.args["graph_save_dir"])
self.yolo = YOLO()
def yolo_boxes_and_scores(feats, anchors, num_classes, input_shape,
image_shape):
'''Process Conv layer output'''
box_xy, box_wh, box_confidence, box_class_probs = YOLO.yolo_head(
feats, anchors, num_classes, input_shape)
boxes = yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape)
# (x, 4)
boxes = np.reshape(boxes, [-1, 4])
box_scores = box_confidence * box_class_probs
# (x, 10)
box_scores = np.reshape(box_scores, [-1, num_classes])
# (x, 4), (x, 10)
return boxes, box_scores
def _main_(args):
config_path = args.conf
weights_path = args.weights
image_path = args.image
with open(config_path) as config_buffer:
config = json.load(config_buffer)
###############################
# Make the model
###############################
yolo = YOLO(architecture=config['model']['architecture'],
input_size=config['model']['input_size'],
labels=config['model']['labels'],
max_box_per_image=config['model']['max_box_per_image'],
anchors=config['model']['anchors'])
###############################
# Load trained weights
###############################
yolo.load_weights(weights_path)
###############################
# Predict bounding boxes
###############################
image = cv2.imread(image_path)
boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])
print len(boxes), 'boxes are found'
cv2.imwrite(image_path[:-4] + '_detected' + image_path[-4:], image)
for epoch in range(epoches):
running_loss = 0.0
for i, (img, truth) in enumerate(dataset):
optimizer.zero_grad()
if cuda:
img, truth = img.cuda(), truth.cuda()
pred = model(img)
loss = criterion(pred, truth)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(loss.item())
if i % 20 == 19:
print("Running loss:%f" % (running_loss / 20))
running_loss = .0
torch.save(net.state_dict(), "yolov4.ptn")
if __name__ == "__main__":
cocoset = YOLOData(img_dir, ann_dir)
coco = DataLoader(cocoset, batch_size=1)
net = YOLO(cfg_file)
net.load_weights("weights/yolov4.weights")
criterion = YOLOLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
train(net, coco, criterion, optimizer)
def yolo4_loss(args):
'''Return yolo4_loss tensor
Parameters
----------
yolo_outputs: list of tensor, the output of yolo_body or tiny_yolo_body
y_true: list of array, the output of preprocess_true_boxes
anchors: array, shape=(N, 2), wh
num_classes: integer
ignore_thresh: float, the iou threshold whether to ignore object confidence loss
Returns
-------
loss: tensor, shape=(1,)
'''
# total_location_loss = 0
# total_confidence_loss = 0
# total_class_loss = 0
loss = 0
anchors, num_classes, ignore_thresh = config.anchors, config.num_classes, config.ignore_thresh
y_pred_base, y_true = args[:3], args[3:]
# 3
num_layers = len(anchors) // 3 # default setting
# (9, 2)
anchors = np.asarray(anchors, dtype=int)
input_shape = K.cast(K.shape(y_pred_base[0])[1:3] * 32, K.dtype(y_true[0]))
grid_shapes = [
K.cast(K.shape(y_pred_base[l])[1:3], K.floatx())
for l in range(num_layers)
]
# N
batch = K.shape(y_pred_base[0])[0]
batch_tensor = K.cast(batch, K.floatx())
for l in range(num_layers):
object_mask = y_true[l][..., 4:5]
true_class_probs = y_true[l][..., 5:]
grid, raw_pred, pred_xy, pred_wh = YOLO().yolo_head(
y_pred_base[l],
anchors[config.anchor_mask[l]],
num_classes,
input_shape,
calc_loss=True)
pred_box = K.concatenate([pred_xy, pred_wh])
# raw_true_xy = y_true[l][..., :2] * grid_shapes[l][::-1] - grid
raw_true_wh = K.log(
y_true[l][..., 2:4] /
(anchors[config.anchor_mask[l]] * input_shape[::-1] + K.epsilon()))
# raw_true_wh = K.switch(object_mask, raw_true_wh, K.zeros_like(raw_true_wh)) # avoid log(0)=-inf
box_loss_scale = 2 - y_true[l][..., 2:3] * y_true[l][..., 3:4]
ignore_mask = tf.TensorArray(K.dtype(y_true[0]),
size=1,
dynamic_size=True)
object_mask_bool = K.cast(object_mask, 'bool')
def loop_body(b, ignore_mask):
true_box = tf.boolean_mask(y_true[l][b, ..., 0:4],
object_mask_bool[b, ..., 0])
iou = utils.iou_cors_index(pred_box[b], true_box)
best_iou = K.max(iou, axis=-1)
ignore_mask = ignore_mask.write(
b, K.cast(best_iou < ignore_thresh, K.dtype(true_box)))
return b + 1, ignore_mask
_, ignore_mask = tf.while_loop(lambda b, *args: b < batch, loop_body,
[0, ignore_mask])
ignore_mask = ignore_mask.stack()
ignore_mask = K.expand_dims(ignore_mask, -1)
confidence_loss = object_mask * K.binary_crossentropy(object_mask, raw_pred[..., 4:5], from_logits=True) + \
(1 - object_mask) * K.binary_crossentropy(object_mask, raw_pred[..., 4:5],
from_logits=True) * ignore_mask
class_loss = object_mask * K.binary_crossentropy(
true_class_probs, raw_pred[..., 5:], from_logits=True)
raw_true_box = y_true[l][..., 0:4]
ciou = box_ciou(pred_box, raw_true_box)
ciou_loss = object_mask * box_loss_scale * ciou
ciou_loss = K.sum(ciou_loss) / batch_tensor
location_loss = ciou_loss
confidence_loss = K.sum(confidence_loss) / batch_tensor
class_loss = K.sum(class_loss) / batch_tensor
loss += location_loss + confidence_loss + class_loss
# total_location_loss += location_loss
# total_confidence_loss += confidence_loss
# total_class_loss += class_loss
loss = K.expand_dims(loss, axis=-1)
return loss
"-1 means cpu.")
parser.add_argument('--num_processes',
type=int,
default=0,
help='number of processes.')
return parser.parse_args().__dict__
if __name__ == '__main__':
args = parse_args()
classes = read_classes(args["class_path"])
color_dict = get_color_dict(classes, args["color_path"])
anchors = read_anchors(
args["anchor_path"]) if args["anchor_path"] else None
model = YOLO(classes,
model_load_path=args["model_load_path"],
anchors=anchors,
device_ids=args["device_ids"])
if args["input_path"].endswith(".jpg") or \
args["input_path"].endswith(".jpeg") or \
args["input_path"].endswith(".png"):
model.detect_image(
args["input_path"],
args["score_threshold"],
args["iou_threshold"],
color_dict,
do_show=args["do_show"],
delay=0,
output_path=args["output_path"] if args["output_path"] else None)
elif args["input_path"].endswith(".mp4"):
model.detect_video(
args["input_path"],
def _main_(args):
config_path = args.conf
with open(config_path) as config_buffer:
config = json.load(config_buffer)
###############################
# Parse the annotations
###############################
# parse annotations of the training set
train_imgs, train_labels = parse_annotation(
config['train']['train_annot_folder'],
config['train']['train_image_folder'], config['model']['labels'])
# parse annotations of the validation set, if any, otherwise split the training set
if os.path.exists(config['valid']['valid_annot_folder']):
valid_imgs, valid_labels = parse_annotation(
config['valid']['valid_annot_folder'],
config['valid']['valid_image_folder'], config['model']['labels'])
else:
train_valid_split = int(0.8 * len(train_imgs))
np.random.shuffle(train_imgs)
valid_imgs = train_imgs[train_valid_split:]
train_imgs = train_imgs[:train_valid_split]
if len(set(config['model']['labels']).intersection(train_labels)) == 0:
print "Labels to be detected are not present in the dataset! Please revise the list of labels in the config.json file!"
return
###############################
# Construct the model
###############################
yolo = YOLO(architecture=config['model']['architecture'],
input_size=config['model']['input_size'],
labels=config['model']['labels'],
max_box_per_image=config['model']['max_box_per_image'],
anchors=config['model']['anchors'])
###############################
# Load the pretrained weights (if any)
###############################
if os.path.exists(config['train']['pretrained_weights']):
yolo.load_weights(config['train']['pretrained_weights'])
###############################
# Start the training process
###############################
yolo.train(train_imgs=train_imgs,
valid_imgs=valid_imgs,
train_times=config['train']['train_times'],
valid_times=config['valid']['valid_times'],
nb_epoch=config['train']['nb_epoch'],
learning_rate=config['train']['learning_rate'],
batch_size=config['train']['batch_size'],
warmup_bs=config['train']['warmup_batches'],
object_scale=config['train']['object_scale'],
no_object_scale=config['train']['no_object_scale'],
coord_scale=config['train']['coord_scale'],
class_scale=config['train']['class_scale'],
debug=config['train']['debug'])
class Classifier:
def __init__(self, args):
print("args = {")
for k in args:
print("\t{} = {}".format(k, args[k]))
print("}")
self.args = args.copy()
if self.args["dataset_path"] != "":
self.data_loader = VOCDataLoader(self.args["dataset_path"],
num_processes=4,
preload=self.args["preload"])
self.classes = read_classes(self.args["class_path"])
self.color_dict = get_color_dict(self.classes, self.args["color_path"])
if self.args["model_save_dir"] != "" and not os.path.exists(
self.args["model_save_dir"]):
os.makedirs(self.args["model_save_dir"])
if self.args["graph_save_dir"] != "" and not os.path.exists(
self.args["graph_save_dir"]):
os.makedirs(self.args["graph_save_dir"])
self.yolo = YOLO()
def run(self):
if self.args["image_detect_path"] != "":
self.yolo.detect_image_and_show(self.args["image_detect_path"],
self.color_dict, 0)
if self.args["video_detect_path"] != "":
self.yolo.detect_video_and_show(
self.args["video_detect_path"],
self.color_dict,
)
if not any([
self.args["do_train"], self.args["do_eval"],
self.args["do_test"]
]):
return None
print('-' * 20 + 'Reading data' + '-' * 20, flush=True)
data_train = self.data_loader.get_data_train(
) if self.args["do_train"] else []
data_eval = self.data_loader.get_data_eval(
) if self.args["do_eval"] else []
data_test = self.data_loader.get_data_test(
) if self.args["do_test"] else []
print('-' * 20 + 'Preprocessing data' + '-' * 20, flush=True)
for data_id in range(len(data_train)):
data_train[data_id][0] = self.yolo.preprocess_image(
*(data_train[data_id]), cvt_RGB=True)[0]
for data_id in range(len(data_eval)):
data_eval[data_id][0] = self.yolo.preprocess_image(
*(data_eval[data_id]), cvt_RGB=True)[0]
for data_id in range(len(data_test)):
data_test[data_id][0] = self.yolo.preprocess_image(
*(data_test[data_id]), cvt_RGB=True)[0]
if self.args["graph_save_dir"] != "":
self.yolo.save_graph(self.args["graph_save_dir"])
for epoch in range(self.args["num_epochs"]):
if self.args["do_train"]:
"""Train"""
print('-' * 20 + 'Training epoch %d' % epoch + '-' * 20,
flush=True)
time.sleep(0.5)
random.shuffle(data_train) # 打乱训练数据
for start in tqdm(range(0, len(data_train),
self.args["train_batch_size"]),
desc='Training batch: '):
end = min(start + self.args["train_batch_size"],
len(data_train))
loss = self.yolo.train(data_train[start:end])
print(loss)
"""Save current model"""
if self.args["model_save_dir"] != "":
self.yolo.save(
os.path.join(
self.args["model_save_dir"],
time.strftime("%Y-%m-%d_%H-%M-%S",
time.localtime()) + "_" +
str(epoch) + ".pth"))
if self.args["do_eval"]:
"""Evaluate"""
print('-' * 20 + 'Evaluating epoch %d' % epoch + '-' * 20,
flush=True)
time.sleep(0.5)
pred_results = []
for start in tqdm(range(0, len(data_eval),
self.args["eval_batch_size"]),
desc='Evaluating batch: '):
end = min(start + self.args["eval_batch_size"],
len(data_eval))
pred_results += self.yolo.predict(
[data[0] for data in data_eval[start:end]],
num_processes=0)
mmAP = self.yolo.get_mmAP(data_eval, pred_results)
print("mmAP =", mmAP)
if not self.args["do_train"]:
break
if self.args["do_test"]:
pass
# """Test"""
# print('-' * 20 + 'Testing epoch %d' % epoch + '-' * 20, flush=True)
# time.sleep(0.1)
# m = metrics.Metrics(self.labels)
# for start in tqdm(range(0, len(self.data_test), self.args.test_batch_size),
# desc='Testing batch: '):
# images = [d[0] for d in self.data_test[start:start + self.args.test_batch_size]]
# actual_labels = [d[1] for d in self.data_test[start:start + self.args.test_batch_size]]
# """forward"""
# batch_images = torch.tensor(images, dtype=torch.float32)
# outputs = self.model(batch_images)
# """update confusion matrix"""
# pred_labels = outputs.softmax(1).argmax(1).tolist()
# m.update(actual_labels, pred_labels)
# """testing"""
# print(m.get_accuracy())
if not self.args["do_train"]:
break
print()
# Get numpy image in opencv format
np_img = preprocess.post_process_image(img1)
# get bounding boxes: xyxy + class + confidence
img_h, img_w, _ = np_img.shape
target_bboxes = preprocess.decode_label(target1, img_h, img_w)
draw_all_bboxes(np_img, target_bboxes, preprocess, gt_color,
"images_transformed/gt_{}.jpg".format(pos))
import os
os._exit(1)
# Enable anomaly detection for debugging purpose
torch.autograd.set_detect_anomaly(True)
model = YOLO((channels, height, width), S, B, C)
if num_gpu > 1:
print("Let's use", num_gpu, "GPUs!")
model = torch.nn.DataParallel(model)
model = model.to(device)
# Init tensorboard for loss visualization
tb = Tensorboard()
loss_func = YoloLoss(S, B, C) # torch.nn.MSELoss(reduction="sum")
optimizer = opt(model.parameters(),
momentum=momentum,
lr=lr,
weight_decay=weight_decay) # momentum=momentum
# Keep the loss of the best model
# Load dataset
test_dataset = Dataset(test_images_dir,
test_files, (height, width),
S,
B,
C,
preprocess,
random_transform=False)
test_generator = data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
# Load model
model = YOLO((channels, height, width), S, B, C)
try:
load_checkpoint(model)
except:
"Try with a dataparallel model"
model = torch.nn.DataParallel(model)
load_checkpoint(model)
model = model.module
model = model.to(device)
loss_func = YoloLoss(S, B, C) # torch.nn.MSELoss(reduction="sum")
start_timestamp = time.time()
# Evaluation
model.eval()
model.backbone = None
torch.cuda.empty_cache()
print("Cache cleaned.")
model_lock.release()
time.sleep(time_gap)
def run_server():
server = pywsgi.WSGIServer(('0.0.0.0', 34560), app, log=app.logger)
server.serve_forever()
# app.run(threaded=False)
if __name__ == '__main__':
classes = read_classes("./data/coco-ch.names")
color_dict = get_color_dict(classes, "./data/colors")
anchors = read_anchors("./data/anchors")
model = YOLO(classes,
model_load_path=model_path,
anchors=anchors,
device_ids="0")
if not os.path.exists(out_dir):
os.makedirs(out_dir)
model_use_time = time.time()
model_lock = threading.Lock()
threading.Thread(target=cuda_memory_control, args=(60, 30)).start()
threading.Thread(target=run_server).start()
print("Server is successfully started.")
label_dict = make_label_dict('data/class_names.txt')
C = len(label_dict)
dataset = YOLODataset('data/train',
S,
B,
C,
label_dict=label_dict,
to_tensor=to_tensor)
dataloader = BaseDataLoader(dataset=dataset, batch_size=args.batch_size)
len_train = len(dataset) * 8 // 10
len_val = len(dataset) - len_train
train_dataloader, val_dataloader = dataloader.split([len_train, len_val])
net = YOLO(S, B, C)
if args.gpu:
net = net.cuda()
net = nn.DataParallel(net)
criterion = Loss(S, B, C)
metric = Loss(S, B, C)
optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),
lr=args.lr,
momentum=0.9)
scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
model = Model(net)
model.compile(optimizer, criterion, metric, scheduler, label_dict)
model.fit(train_dataloader=train_dataloader,
class Yolo4(object):
def __init__(self, model_path, weights_path, gpu_num=1):
if not self.check_model(model_path):
self.check_weights(weights_path)
self.score = config.score
self.iou = config.iou
self.weights_path = weights_path
self.model_path = model_path
self.gpu_num = gpu_num
self.colors = utils_image.get_random_colors(
len(config.classes_names))
self.yolo4_model = YOLO(config.image_input_shape)()
self.convertor()
print('Converting finished !')
@staticmethod
def check_model(model_path: str):
if os.path.exists(model_path):
print(model_path + ' exists, stop converting!')
return 1
return
@staticmethod
def check_weights(weights_path: str):
if weights_path.startswith('/'):
root, file_name = os.path.split(model_path)
elif weights_path.find('/') > 0:
root, file_name = weights_path.split('/')
else:
root = ''
file_name = weights_path
file_pre_name = file_name.split('.')[0]
file_lock_name = root + '/.' + file_pre_name + '.lock'
if os.path.exists(file_lock_name):
return 'Weights file has been transformed yet!'
else:
with open(weights_path, 'rb') as f:
data = f.read()
f.close()
with open(weights_path, 'wb') as f:
f.write(data[20:] + data[:20])
f.close()
with open(file_lock_name, 'w') as f:
f.write('')
f.close()
return 'Weights file has been transformed successfully!'
def convertor(self):
weights_file = open(self.weights_path, 'rb')
convs_to_load = []
bns_to_load = []
for i in range(len(self.yolo4_model.layers)):
layer_name = self.yolo4_model.layers[i].name
layer_name = utils.tf_layer_name_compat(layer_name)
if layer_name.startswith('conv2d_'):
convs_to_load.append((int(layer_name[7:]), i))
if layer_name.startswith('batch_normalization_'):
bns_to_load.append((int(layer_name[20:]), i))
convs_sorted = sorted(convs_to_load, key=itemgetter(0))
bns_sorted = sorted(bns_to_load, key=itemgetter(0))
bn_index = 0
for i in range(len(convs_sorted)):
# print('Converting ', i)
if i == 93 or i == 101 or i == 109:
# no bn, with bias
weights_shape = self.yolo4_model.layers[
convs_sorted[i][1]].get_weights()[0].shape
bias_shape = self.yolo4_model.layers[
convs_sorted[i][1]].get_weights()[0].shape[3]
filters = bias_shape
size = weights_shape[0]
darknet_w_shape = (filters, weights_shape[2], size, size)
weights_size = np.product(weights_shape)
# exit()
conv_bias = np.ndarray(shape=(filters, ),
dtype='float32',
buffer=weights_file.read(filters * 4))
conv_weights = np.ndarray(shape=darknet_w_shape,
dtype='float32',
buffer=weights_file.read(
weights_size * 4))
conv_weights = np.transpose(conv_weights, [2, 3, 1, 0])
self.yolo4_model.layers[convs_sorted[i][1]].set_weights(
[conv_weights, conv_bias])
else:
# with bn, no bias
weights_shape = self.yolo4_model.layers[
convs_sorted[i][1]].get_weights()[0].shape
size = weights_shape[0]
bn_shape = self.yolo4_model.layers[bns_sorted[bn_index]
[1]].get_weights()[0].shape
filters = bn_shape[0]
darknet_w_shape = (filters, weights_shape[2], size, size)
weights_size = np.product(weights_shape)
conv_bias = np.ndarray(shape=(filters, ),
dtype='float32',
buffer=weights_file.read(filters * 4))
bn_weights = np.ndarray(shape=(3, filters),
dtype='float32',
buffer=weights_file.read(filters * 12))
bn_weight_list = [
bn_weights[0], conv_bias, bn_weights[1], bn_weights[2]
]
self.yolo4_model.layers[bns_sorted[bn_index][1]].set_weights(
bn_weight_list)
conv_weights = np.ndarray(shape=darknet_w_shape,
dtype='float32',
buffer=weights_file.read(
weights_size * 4))
conv_weights = np.transpose(conv_weights, [2, 3, 1, 0])
self.yolo4_model.layers[convs_sorted[i][1]].set_weights(
[conv_weights])
bn_index += 1
weights_file.close()
self.yolo4_model.save(self.model_path)
def yolo4_loss(args,
anchors,
num_classes,
ignore_thresh=.5,
label_smoothing=0,
use_focal_loss=False,
use_focal_obj_loss=False,
use_softmax_loss=False,
use_giou_loss=False,
use_diou_loss=False):
'''Return yolo4_loss tensor
Parameters
----------
yolo_outputs: list of tensor, the output of yolo_body or tiny_yolo_body
y_true: list of array, the output of preprocess_true_boxes
anchors: array, shape=(N, 2), wh
num_classes: integer
ignore_thresh: float, the iou threshold whether to ignore object confidence loss
Returns
-------
loss: tensor, shape=(1,)
'''
total_location_loss = 0
total_confidence_loss = 0
total_class_loss = 0
y_pred_base, y_true = args[:3], args[3:]
# 3
num_layers = len(anchors) // 3 # default setting
# (9, 2)
anchors = np.asarray(anchors, dtype=int)
# (416, 416)
input_shape = K.cast(K.shape(y_pred_base[0])[1:3] * 32, K.dtype(y_true[0]))
# [(13, 13), (26, 26), (52, 52)]
grid_shapes = [
K.cast(K.shape(y_pred_base[l])[1:3], K.floatx())
for l in range(num_layers)
]
loss = 0
# N
batch = K.shape(y_pred_base[0])[0] # batch size, tensor
batch_tensor = K.cast(batch, K.floatx())
for l in range(num_layers):
object_mask = y_true[l][..., 4:5]
true_class_probs = y_true[l][..., 5:]
grid, raw_pred, pred_xy, pred_wh = YOLO.yolo_head(
y_pred_base[l],
anchors[config.anchor_mask[l]],
num_classes,
input_shape,
calc_loss=True)
pred_box = K.concatenate([pred_xy, pred_wh])
# Darknet raw box to calculate loss.
raw_true_xy = y_true[l][..., :2] * grid_shapes[l][::-1] - grid
raw_true_wh = K.log(y_true[l][..., 2:4] /
anchors[config.anchor_mask[l]] * input_shape[::-1])
raw_true_wh = K.switch(object_mask, raw_true_wh,
K.zeros_like(raw_true_wh)) # avoid log(0)=-inf
box_loss_scale = 2 - y_true[l][..., 2:3] * y_true[l][..., 3:4]
# Find ignore mask, iterate over each of batch.
ignore_mask = tf.TensorArray(K.dtype(y_true[0]),
size=1,
dynamic_size=True)
object_mask_bool = K.cast(object_mask, 'bool')
def loop_body(b, ignore_mask):
true_box = tf.boolean_mask(y_true[l][b, ..., 0:4],
object_mask_bool[b, ..., 0])
iou = utils.iou_cors_index(pred_box[b], true_box)
best_iou = K.max(iou, axis=-1)
ignore_mask = ignore_mask.write(
b, K.cast(best_iou < ignore_thresh, K.dtype(true_box)))
return b + 1, ignore_mask
# _, ignore_mask = tf.while_loop(lambda b, *args: b < m, loop_body, [0, ignore_mask])
_, ignore_mask = tf.while_loop(lambda b, *args: b < batch, loop_body,
[0, ignore_mask])
ignore_mask = ignore_mask.stack()
ignore_mask = K.expand_dims(ignore_mask, -1)
confidence_loss = object_mask * K.binary_crossentropy(object_mask, raw_pred[..., 4:5], from_logits=True) + \
(1 - object_mask) * K.binary_crossentropy(object_mask, raw_pred[..., 4:5],
from_logits=True) * ignore_mask
class_loss = object_mask * K.binary_crossentropy(
true_class_probs, raw_pred[..., 5:], from_logits=True)
# Calculate DIoU loss as location loss
raw_true_box = y_true[l][..., 0:4]
diou = box_diou(pred_box, raw_true_box)
diou_loss = object_mask * box_loss_scale * (1 - diou)
diou_loss = K.sum(diou_loss) / batch_tensor
location_loss = diou_loss
confidence_loss = K.sum(confidence_loss) / batch_tensor
class_loss = K.sum(class_loss) / batch_tensor
loss += location_loss + confidence_loss + class_loss
total_location_loss += location_loss
total_confidence_loss += confidence_loss
total_class_loss += class_loss
# Fit for tf 2.0.0 loss shape
loss = K.expand_dims(loss, axis=-1)
return loss # , total_location_loss, total_confidence_loss, total_class_loss