def loss_boxes(p_bbox, p_class, t_bbox, t_class, t_indices, p_indices, t_selector, p_selector):
#print("------")
p_bbox = tf.gather(p_bbox, p_indices)
t_bbox = tf.gather(t_bbox, t_indices)
p_bbox_xy = bbox.xcycwh_to_xy_min_xy_max(p_bbox)
t_bbox_xy = bbox.xcycwh_to_xy_min_xy_max(t_bbox)
l1_loss = tf.abs(p_bbox-t_bbox)
l1_loss = tf.reduce_sum(l1_loss) / tf.cast(tf.shape(p_bbox)[0], tf.float32)
iou, union = bbox.jaccard(p_bbox_xy, t_bbox_xy, return_union=True)
_p_bbox_xy, _t_bbox_xy = bbox.merge(p_bbox_xy, t_bbox_xy)
top_left = tf.math.minimum(_p_bbox_xy[:,:,:2], _t_bbox_xy[:,:,:2])
bottom_right = tf.math.maximum(_p_bbox_xy[:,:,2:], _t_bbox_xy[:,:,2:])
size = tf.nn.relu(bottom_right - top_left)
area = size[:,:,0] * size[:,:,1]
giou = (iou - (area - union) / area)
loss_giou = 1 - tf.linalg.diag_part(giou)
loss_giou = tf.reduce_sum(loss_giou) / tf.cast(tf.shape(p_bbox)[0], tf.float32)
return loss_giou, l1_loss
def hungarian_matching(t_bbox,
t_class,
p_bbox,
p_class,
lwts,
slice_preds=True) -> tuple:
if slice_preds:
n_objs = tf.cast(
t_bbox[0][0], tf.int32
) # this takes the num_bboxes from the padded header row, containing (say) [400, 0, 0, 0], hence n_objs would be 400
t_bbox = tf.slice(
t_bbox, [1, 0],
[n_objs, 4
]) # slice begins at first row to get rid of the header row
t_class = tf.slice(t_class, [1, 0], [n_objs, -1])
# t_class = tf.squeeze(t_class, axis=-1)
# Convert from [xc, yc, w, h] to [xmin, ymin, xmax, ymax]
p_bbox_xy = xcycwh_to_xy_min_xy_max(p_bbox)
t_bbox_xy = xcycwh_to_xy_min_xy_max(t_bbox)
# t_bbox_xy = tf.cast(bbox.xcycwh_to_xy_min_xy_max(t_bbox), tf.float64)
# Classification cost for the Hungarian algorithom
# On each prediction. We select the prob of the expected class
sigmoid = tf.math.sigmoid(
p_class) # this is because detr has exclusive classes, we don't!
_p_class, _t_class = merge(sigmoid, t_class)
cost_class = tf.norm(_p_class - _t_class, ord=1, axis=-1)
## we could introduce different class weightings by writing a custom
## norm fucntion above, multipling classes separately before summing
# L1 cost for the hungarian algorithm
_p_bbox, _t_bbox = merge(p_bbox, t_bbox)
cost_bbox = tf.norm(_p_bbox - _t_bbox, ord=1, axis=-1)
# Generalized IOU
iou, union = jaccard(p_bbox_xy, t_bbox_xy, return_union=True)
_p_bbox_xy, _t_bbox_xy = merge(p_bbox_xy, t_bbox_xy)
top_left = tf.math.minimum(_p_bbox_xy[:, :, :2], _t_bbox_xy[:, :, :2])
bottom_right = tf.math.maximum(_p_bbox_xy[:, :, 2:], _t_bbox_xy[:, :, 2:])
size = tf.nn.relu(bottom_right - top_left)
area = size[:, :, 0] * size[:, :, 1]
cost_giou = -(iou - (area - union) / area)
# Final hungarian cost matrix
# loss_weights = [1, 2, 5] # label_cost, giou_loss, l1_loss
cost_matrix = lwts[2] * cost_bbox + lwts[0] * cost_class + lwts[
1] * cost_giou
selectors = tf.numpy_function(np_tf_linear_sum_assignment, [cost_matrix],
[tf.int64, tf.int64, tf.bool, tf.bool])
target_indices = selectors[0]
pred_indices = selectors[1]
target_selector = selectors[2]
pred_selector = selectors[3]
return pred_indices, target_indices, pred_selector, target_selector, t_bbox, t_class
def get_model_inference(m_outputs: dict,
background_class,
bbox_format="xy_center"):
predicted_bbox = m_outputs["pred_boxes"][0]
predicted_labels = m_outputs["pred_logits"][0]
softmax = tf.nn.softmax(predicted_labels)
predicted_scores = tf.reduce_max(softmax, axis=-1)
predicted_labels = tf.argmax(softmax, axis=-1)
indices = tf.where(predicted_labels != background_class)
indices = tf.squeeze(indices, axis=-1)
predicted_scores = tf.gather(predicted_scores, indices)
predicted_labels = tf.gather(predicted_labels, indices)
predicted_bbox = tf.gather(predicted_bbox, indices)
if bbox_format == "xy_center":
predicted_bbox = predicted_bbox
elif bbox_format == "xyxy":
predicted_bbox = bbox.xcycwh_to_xy_min_xy_max(predicted_bbox)
elif bbox_format == "yxyx":
predicted_bbox = bbox.xcycwh_to_yx_min_yx_max(predicted_bbox)
else:
raise NotImplementedError()
return predicted_bbox, predicted_labels, predicted_scores
def hungarian_matching(t_bbox, t_class, p_bbox, p_class, fcost_class=1, fcost_bbox=5, fcost_giou=2, slice_preds=True) -> tuple:
if slice_preds:
size = tf.cast(t_bbox[0][0], tf.int32)
t_bbox = tf.slice(t_bbox, [1, 0], [size, 4])
t_class = tf.slice(t_class, [1, 0], [size, -1])
t_class = tf.squeeze(t_class, axis=-1)
# Convert frpm [xc, yc, w, h] to [xmin, ymin, xmax, ymax]
p_bbox_xy = bbox.xcycwh_to_xy_min_xy_max(p_bbox)
t_bbox_xy = bbox.xcycwh_to_xy_min_xy_max(t_bbox)
softmax = tf.nn.softmax(p_class)
# Classification cost for the Hungarian algorithom
# On each prediction. We select the prob of the expected class
cost_class = -tf.gather(softmax, t_class, axis=1)
# L1 cost for the hungarian algorithm
_p_bbox, _t_bbox = bbox.merge(p_bbox, t_bbox)
cost_bbox = tf.norm(_p_bbox - _t_bbox, ord=1, axis=-1)
# Generalized IOU
iou, union = bbox.jaccard(p_bbox_xy, t_bbox_xy, return_union=True)
_p_bbox_xy, _t_bbox_xy = bbox.merge(p_bbox_xy, t_bbox_xy)
top_left = tf.math.minimum(_p_bbox_xy[:,:,:2], _t_bbox_xy[:,:,:2])
bottom_right = tf.math.maximum(_p_bbox_xy[:,:,2:], _t_bbox_xy[:,:,2:])
size = tf.nn.relu(bottom_right - top_left)
area = size[:,:,0] * size[:,:,1]
cost_giou = -(iou - (area - union) / area)
# Final hungarian cost matrix
cost_matrix = fcost_bbox * cost_bbox + fcost_class * cost_class + fcost_giou * cost_giou
selectors = tf.numpy_function(np_tf_linear_sum_assignment, [cost_matrix], [tf.int64, tf.int64, tf.bool, tf.bool] )
target_indices = selectors[0]
pred_indices = selectors[1]
target_selector = selectors[2]
pred_selector = selectors[3]
return pred_indices, target_indices, pred_selector, target_selector, t_bbox, t_class
def tf_send_batch_log_to_wandb(images,
target_bbox,
target_class,
m_outputs: dict,
config,
class_name=[],
step=None,
prefix=""):
# Warning: In graph mode, this class is init only once. In eager mode, this class is init at each step.
img_sender = WandbSender()
predicted_bbox = m_outputs["pred_boxes"]
for b in range(predicted_bbox.shape[0]):
# Select within the batch the elements at indice b
image = images[b]
elem_m_outputs = {
key: m_outputs[key][b:b + 1] if
(m_outputs[key] is not None
and not isinstance(m_outputs[key], list)) else m_outputs[key]
for key in m_outputs
}
# Target
t_bbox, t_class = target_bbox[b], target_class[b]
if not RAGGED:
size = tf.cast(t_bbox[0][0], tf.int32)
t_bbox = tf.slice(t_bbox, [1, 0], [size, 4])
t_bbox = bbox.xcycwh_to_xy_min_xy_max(t_bbox)
t_class = tf.slice(t_class, [1, 0], [size, -1])
t_class = tf.squeeze(t_class, axis=-1)
# Predictions
predicted_bbox, predicted_labels, predicted_scores = get_model_inference(
elem_m_outputs, config.background_class, bbox_format="xyxy")
np_func_params = {
"image": image,
"p_bbox": np.array(predicted_bbox),
"p_scores": np.array(predicted_scores),
"t_bbox": np.array(t_bbox),
"p_labels": np.array(predicted_labels),
"t_labels": np.array(t_class),
"class_name": class_name
}
img_sender.gather_inference(**np_func_params)
img_sender.send(step=step, prefix=prefix)
model.save_weights(weights_name_next)
# test on a batch
logits = model(images)
losses = batched_loss_closure([t_bbox, t_class], logits, loss_weights,
class_weights)
total_loss = tf.reduce_mean(
tf.gather(losses, [tf.shape(losses)[1] - 1], axis=1))
#total_loss = calc_loss_alone([t_bbox, t_class], logits)
from bbox import xcycwh_to_xy_min_xy_max
import cv2
bboxes = [
xcycwh_to_xy_min_xy_max(logits[0][b, :, :])
for b in range(logits[0].shape[0])
]
p_is = tf.math.sigmoid(logits[1])
norm_std = train_datagen.norm_std
norm_mean = train_datagen.norm_mean
for ba in range(images.shape[0]):
img = (255 * (images[ba] * norm_std + norm_mean)).astype(np.uint8)
bbx_ij = bboxes[ba].numpy()
p_i = p_is[ba, :, :].numpy()
for j in range(p_i.shape[0]):
bbx_j = img.shape[0] * bbx_ij[j]
pi_j = np.round(p_i[j], 2)
x = int(bbx_j[0])
y = int(bbx_j[1])