diff --git a/config/yoeo.cfg b/config/yoeo.cfg index 100e9d1..55464fe 100644 --- a/config/yoeo.cfg +++ b/config/yoeo.cfg @@ -4,7 +4,7 @@ #subdivisions=1 # Training batch=64 -subdivisions=8 +subdivisions=4 width=416 height=416 channels=3 @@ -19,8 +19,8 @@ learning_rate=0.001 burn_in=100 max_batches = 4000 policy=steps -steps=50000,60000 -scales=.1,.1 +steps=80000,250000 +scales=.1,.2 #### # Like YOEO rev 2 but with deeper skip connections @@ -213,15 +213,15 @@ activation=leaky size=1 stride=1 pad=1 -filters=24 +filters=12 activation=linear [yolo] -mask = 3,4,5 -anchors = 17, 32, 13,171, 37, 67, 30,224, 69,112, 116,212 -classes=3 +mask = 0 +anchors = 100, 100 +classes=4 num=6 jitter=.3 ignore_thresh = .7 @@ -257,13 +257,13 @@ activation=leaky size=1 stride=1 pad=1 -filters=24 +filters=12 activation=linear [yolo] -mask = 0,1,2 -anchors = 17, 32, 13,171, 37, 67, 30,224, 69,112, 116,212 -classes=3 +mask = 0 +anchors = 100, 100 +classes=4 num=6 jitter=.3 ignore_thresh = .7 diff --git a/yoeo/detect.py b/yoeo/detect.py index f6f043a..4be7e86 100755 --- a/yoeo/detect.py +++ b/yoeo/detect.py @@ -71,10 +71,10 @@ def detect_directory(model_path, weights_path, img_path, class_config: ClassConf print(f"---- Detections were saved to: '{output_path}' ----") -def detect_image(model, - image: np.ndarray, - img_size: int = 416, - conf_thres: float = 0.5, +def detect_image(model, + image: np.ndarray, + img_size: int = 416, + conf_thres: float = 0.5, nms_thres: float = 0.5, group_config: Optional[GroupConfig] = None ): @@ -97,7 +97,7 @@ def detect_image(model, :rtype: nd.array, nd.array """ model.eval() # Set model to evaluation mode - + # Configure input input_img = transforms.Compose([ DEFAULT_TRANSFORMS, @@ -113,9 +113,9 @@ def detect_image(model, with torch.no_grad(): detections, segmentations = model(input_img) detections = non_max_suppression( - prediction=detections, - conf_thres=conf_thres, - iou_thres=nms_thres, + prediction=detections, + conf_thres=conf_thres, + iou_thres=nms_thres, group_config=group_config ) detections = rescale_boxes(detections[0], img_size, image.shape[0:2]) @@ -124,9 +124,9 @@ def detect_image(model, def detect(model, - dataloader: DataLoader, - output_path: str, - conf_thres: float = 0.5, + dataloader: DataLoader, + output_path: str, + conf_thres: float = 0.5, nms_thres: float = 0.5, group_config: Optional[GroupConfig] = None ): @@ -169,9 +169,9 @@ def detect(model, with torch.no_grad(): detections, segmentations = model(input_imgs) detections = non_max_suppression( - prediction=detections, - conf_thres=conf_thres, - iou_thres=nms_thres, + prediction=detections, + conf_thres=conf_thres, + iou_thres=nms_thres, group_config=group_config ) @@ -254,7 +254,7 @@ def _draw_and_save_output_image(image_path, detections, seg, img_size, output_pa # Bounding-box colors cmap = plt.get_cmap("tab20b") colors = [cmap(i) for i in np.linspace(0, 1, len(classes))] - for x1, y1, x2, y2, conf, cls_pred in detections: + for x1, y1, x2, y2, conf, cls_pred, b, bx, by in detections: print(f"\t+ Label: {classes[int(cls_pred)]} | Confidence: {conf.item():0.4f}") @@ -265,8 +265,12 @@ def _draw_and_save_output_image(image_path, detections, seg, img_size, output_pa bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=colors[int(cls_pred)], facecolor="none") # Add the bbox to the plot ax.add_patch(bbox) + + # Add red dot for base_footprint if it exists + if b > 0.5: + ax.plot(bx, by, 'ro') + # Add label - """ plt.text( x1, y1, @@ -274,7 +278,6 @@ def _draw_and_save_output_image(image_path, detections, seg, img_size, output_pa color="white", verticalalignment="top", bbox={"color": colors[int(cls_pred)], "pad": 0}) - """ # Save generated image with detections plt.axis("off") diff --git a/yoeo/models.py b/yoeo/models.py index 784627c..7b5d65f 100644 --- a/yoeo/models.py +++ b/yoeo/models.py @@ -135,7 +135,7 @@ def __init__(self, anchors, num_classes): self.num_classes = num_classes self.mse_loss = nn.MSELoss() self.bce_loss = nn.BCELoss() - self.no = num_classes + 5 # number of outputs per anchor + self.no = num_classes + 5 + 3 # TODO make basefootprint param # number of outputs per anchor self.grid = torch.zeros(1) # TODO anchors = torch.tensor(list(chain(*anchors))).float().view(-1, 2) @@ -157,7 +157,9 @@ def forward(self, x, img_size): x = torch.cat([ (x[..., 0:2].sigmoid() + self.grid) * stride, # xy torch.exp(x[..., 2:4]) * self.anchor_grid, # wh - x[..., 4:].sigmoid(), + x[..., 4:-3].sigmoid(), + x[..., -3, None].sigmoid(), # basefootprint confidence + (x[..., -2:].tan() + self.grid) * stride # basefoot print xy ], axis=4).view(bs, -1, self.no) return x diff --git a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py index 443e4b5..420a33a 100755 --- a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py +++ b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py @@ -30,6 +30,7 @@ def range_limited_float_type_0_to_1(arg): parser.add_argument("--skip-concealed", action="store_true", help="Skip concealed labels") parser.add_argument("--skip-classes", nargs="+", default=[], help="These bounding box classes will be skipped") parser.add_argument("--robots-with-team-colors", action="store_true", help="The robot class will be subdivided into subclasses, one for each team color (currently either 'blue', 'red' or 'unknown').") +parser.add_argument("--add-base-footprint", action="store_true", help="Add base_footprint to bounding box annotations") args = parser.parse_args() # Available classes for YOEO @@ -54,6 +55,9 @@ def range_limited_float_type_0_to_1(arg): dataset_collection_dir = args.dataset_collection_dir destination_dir = args.dataset_collection_dir +# Keep track of all classes with base_footprint annotations +classes_with_base_footprint = set() + # Overwrite defaults, if destination path is given if args.destination_dir: create_symlinks = True @@ -138,10 +142,10 @@ def range_limited_float_type_0_to_1(arg): (args.skip_concealed and annotation.get('concealed', False))): continue elif class_name in CLASSES['bb_classes']: # Handle bounding boxes - min_x = min(map(lambda x: x[0], annotation['vector'])) - max_x = max(map(lambda x: x[0], annotation['vector'])) - min_y = min(map(lambda x: x[1], annotation['vector'])) - max_y = max(map(lambda x: x[1], annotation['vector'])) + min_x = min(pt[0] for pt in annotation['vector']) + max_x = max(pt[0] for pt in annotation['vector']) + min_y = min(pt[1] for pt in annotation['vector']) + max_y = max(pt[1] for pt in annotation['vector']) annotation_width = max_x - min_x annotation_height = max_y - min_y @@ -154,8 +158,27 @@ def range_limited_float_type_0_to_1(arg): relative_center_y = center_y / img_height # Derive classID from index in predefined classes - classID = CLASSES['bb_classes'].index(class_name) - annotations.append(f"{classID} {relative_center_x} {relative_center_y} {relative_annotation_width} {relative_annotation_height}") + classID = CLASSES['bb_classes'].index(class_name) + + # Serialize bounding box + serialized_bounding_box = f"{classID} {relative_center_x} {relative_center_y} {relative_annotation_width} {relative_annotation_height}" + + # Extract other properties + if args.add_base_footprint: + # Default value if key point is not present + base_footprint_x, base_footprint_y = -1, -1 + # Check if the annotation exists (some classes do not have base_footprint) + if 'base_footprint' in annotation: + classes_with_base_footprint.add(class_name) + # Check if the base_footprint is not None (negative label) + if annotation['base_footprint']: + base_footprint_x = annotation['base_footprint'][0] / img_width + base_footprint_y = annotation['base_footprint'][1] / img_height + # Serialize base_footprint key point + serialized_bounding_box += f" {base_footprint_x} {base_footprint_y}" + + # Add row to file + annotations.append(serialized_bounding_box) else: print(f"The annotation type '{class_name}' is not supported. Image: '{img_name_with_extension}'") @@ -180,8 +203,14 @@ def range_limited_float_type_0_to_1(arg): names_path = os.path.join(destination_dir, "yoeo_names.yaml") names = { 'detection': CLASSES['bb_classes'], - 'segmentation': CLASSES["segmentation_classes"], + 'segmentation': CLASSES["segmentation_classes"] } + +# Add base_footprint classes to names file +if args.add_base_footprint: + names['base_footprint'] = list(classes_with_base_footprint) + +# Save names to yaml file with open(names_path, "w") as names_file: yaml.dump(names, names_file) diff --git a/yoeo/test.py b/yoeo/test.py index 224338d..12d9105 100755 --- a/yoeo/test.py +++ b/yoeo/test.py @@ -50,12 +50,11 @@ def evaluate_model_file(model_path, weights_path, img_path, class_config, batch_ :type nms_thres: float, optional :param verbose: If True, prints stats of model, defaults to True :type verbose: bool, optional - :return: Returns precision, recall, AP, f1, ap_class """ dataloader = _create_validation_data_loader( img_path, batch_size, img_size, n_cpu) model = load_model(model_path, weights_path) - metrics_output, seg_class_ious, secondary_metric = _evaluate( + metrics_output, seg_class_ious, secondary_metric, base_footprint_metric = _evaluate( model, dataloader, class_config, @@ -64,13 +63,14 @@ def evaluate_model_file(model_path, weights_path, img_path, class_config, batch_ conf_thres, nms_thres, verbose) - return metrics_output, seg_class_ious, secondary_metric + return metrics_output, seg_class_ious, secondary_metric, base_footprint_metric -def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]], - seg_class_ious: List[np.float64], - secondary_metric: Optional[Metric], - class_config: ClassConfig, +def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]], + seg_class_ious: List[np.float64], + secondary_metric: Optional[Metric], + base_footprint_metrics: dict[str, float], + class_config: ClassConfig, verbose: bool ): # Print detection statistics @@ -80,7 +80,7 @@ def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]], if verbose: # Prints class AP and mean AP ap_table = [["Index", "Class", "AP"]] - class_names = class_config.get_squeezed_det_class_names() + class_names = class_config.get_ungrouped_det_class_names() for i, c in enumerate(ap_class): ap_table += [[c, class_names[c], "%.5f" % AP[i]]] print(AsciiTable(ap_table).table) @@ -95,7 +95,7 @@ def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]], if verbose: classes = class_config.get_group_class_names() mbACC_per_class = [secondary_metric.bACC(i) for i in range(len(classes))] - + sec_table = [["Index", "Class", "bACC"]] for i, c in enumerate(classes): sec_table += [[i, c, "%.5f" % mbACC_per_class[i]]] @@ -116,6 +116,13 @@ def print_eval_stats(metrics_output: Optional[Tuple[np.ndarray]], mean_seg_class_ious = np.array(seg_class_ious).mean() print(f"----Average IoU {mean_seg_class_ious:.5f} ----") + print("#### Base Footprint ####") + # Print base footprint statistics + base_footprint_table = [["Metric", "Value"]] + for metric, value in base_footprint_metrics.items(): + base_footprint_table += [[metric, "%.5f" % value]] + print(AsciiTable(base_footprint_table).table) + def _evaluate(model, dataloader, class_config, img_size, iou_thres, conf_thres, nms_thres, verbose): """Evaluate model on validation dataset. @@ -136,7 +143,7 @@ def _evaluate(model, dataloader, class_config, img_size, iou_thres, conf_thres, :type nms_thres: float :param verbose: If True, prints stats of model :type verbose: bool - :return: Returns precision, recall, AP, f1, ap_class + :return: Returns """ model.eval() # Set model to evaluation mode @@ -157,13 +164,13 @@ def _evaluate(model, dataloader, class_config, img_size, iou_thres, conf_thres, # Extract labels labels += bb_targets[:, 1].tolist() - # If a subset of the detection classes should be grouped into one class for non-maximum suppression and the + # If a subset of the detection classes should be grouped into one class for non-maximum suppression and the # subsequent AP-computation, we need to group those class labels here. if class_config.classes_should_be_grouped(): labels = class_config.group(labels) # Rescale target - bb_targets[:, 2:] = xywh2xyxy(bb_targets[:, 2:]) + bb_targets[:, 2:-2] = xywh2xyxy(bb_targets[:, 2:-2]) bb_targets[:, 2:] *= img_size imgs = Variable(imgs.type(Tensor), requires_grad=False) @@ -180,9 +187,9 @@ def _evaluate(model, dataloader, class_config, img_size, iou_thres, conf_thres, ) sample_stat, secondary_stat = get_batch_statistics( - yolo_outputs, - bb_targets, - iou_threshold=iou_thres, + yolo_outputs, + bb_targets, + iou_threshold=iou_thres, group_config=class_config.get_group_config() ) @@ -200,12 +207,28 @@ def _evaluate(model, dataloader, class_config, img_size, iou_thres, conf_thres, print(f"Times: Mean {1 / np.array(times).mean()}fps | Std: {np.array(times).std()} ms") # Concatenate sample statistics - true_positives, pred_scores, pred_labels = [ + true_positives, pred_scores, pred_labels, base_bootprint_visibility_tp, base_bootprint_visibility_fn, base_bootprint_visibility_fp, base_bootprint_visibility_tn, base_footprint_point_distances = [ np.concatenate(x, 0) for x in list(zip(*sample_metrics))] - + + # Calculate base footprint metrics + base_bootprint_visibility_fn = base_bootprint_visibility_fn.sum() + base_bootprint_visibility_fp = base_bootprint_visibility_fp.sum() + base_bootprint_visibility_tp = base_bootprint_visibility_tp.sum() + base_bootprint_visibility_tn = base_bootprint_visibility_tn.sum() + base_footprint_visibility_precision = base_bootprint_visibility_tp / (base_bootprint_visibility_tp + base_bootprint_visibility_fp + np.finfo(float).eps) + base_footprint_visibility_recall = base_bootprint_visibility_tp / (base_bootprint_visibility_tp + base_bootprint_visibility_fn + np.finfo(float).eps) + base_footprint_metrics = { + "precision": base_footprint_visibility_precision, + "recall": base_footprint_visibility_recall, + "f1": 2 * base_footprint_visibility_precision * base_footprint_visibility_recall / (base_footprint_visibility_precision + base_footprint_visibility_recall + np.finfo(float).eps), + 'keypoint_distance': base_footprint_point_distances.mean() + } + + # Calculate yolo metrics yolo_metrics_output = ap_per_class( true_positives, pred_scores, pred_labels, labels) + # Calculate segmentation metrics def seg_iou_mean_without_nan(seg_iou: List[float]) -> np.ndarray: """This helper function is needed to remove cases, where the segmentation IOU is NaN. This is the case, if a whole batch does not contain any pixels of a segmentation class. @@ -218,9 +241,9 @@ def seg_iou_mean_without_nan(seg_iou: List[float]) -> np.ndarray: seg_class_ious = [seg_iou_mean_without_nan(class_ious) for class_ious in list(zip(*seg_ious))] - print_eval_stats(yolo_metrics_output, seg_class_ious, secondary_metric, class_config, verbose) + print_eval_stats(yolo_metrics_output, seg_class_ious, secondary_metric, base_footprint_metrics, class_config, verbose) - return yolo_metrics_output, seg_class_ious, secondary_metric + return yolo_metrics_output, seg_class_ious, secondary_metric, base_footprint_metrics def _create_validation_data_loader(img_path, batch_size, img_size, n_cpu): diff --git a/yoeo/train.py b/yoeo/train.py index 6f582fe..5d466c9 100755 --- a/yoeo/train.py +++ b/yoeo/train.py @@ -167,9 +167,9 @@ def run(): for batch_i, (_, imgs, bb_targets, mask_targets) in enumerate(tqdm.tqdm(dataloader, desc=f"Training Epoch {epoch}")): batches_done = len(dataloader) * epoch + batch_i - imgs = Variable(imgs.to(device, non_blocking=True)) - bb_targets = Variable(bb_targets.to(device), requires_grad=False) - mask_targets = Variable(mask_targets.to(device=device), requires_grad=False) + imgs = imgs.to(device, non_blocking=True) + bb_targets = bb_targets.to(device) + mask_targets = mask_targets.to(device=device) outputs = model(imgs) @@ -214,8 +214,9 @@ def run(): ["IoU loss", float(loss_components[0])], ["Object loss", float(loss_components[1])], ["Class loss", float(loss_components[2])], - ["Segmentation loss", float(loss_components[3])], - ["Loss", float(loss_components[4])], + ["Base footprint point loss", float(loss_components[3])], + ["Base footprint visibility loss", float(loss_components[4])], + ["Segmentation loss", float(loss_components[5])], ["Batch loss", to_cpu(loss).item()], ]).table) @@ -224,7 +225,9 @@ def run(): ("train/iou_loss", float(loss_components[0])), ("train/obj_loss", float(loss_components[1])), ("train/class_loss", float(loss_components[2])), - ("train/seg_loss", float(loss_components[3])), + ("train/base_footprint_point_loss", float(loss_components[3])), + ("train/base_footprint_visibility_loss", float(loss_components[4])), + ("train/seg_loss", float(loss_components[5])), ("train/loss", to_cpu(loss).item())] logger.list_of_scalars_summary(tensorboard_log, batches_done) @@ -266,8 +269,13 @@ def run(): ("validation/recall", recall.mean()), ("validation/mAP", AP.mean()), ("validation/f1", f1.mean()), - ("validation/seg_iou", np.array(seg_class_ious).mean())] - + ("validation/seg_iou", np.array(seg_class_ious).mean()), + ("validation/base_footprint_dist", metrics_output[3]['keypoint_distance']), + ("validation/base_footprint_visibility_precision", metrics_output[3]['precision']), + ("validation/base_footprint_visibility_recall", metrics_output[3]['recall']), + ("validation/base_footprint_visibility_f1", metrics_output[3]['f1']) + ] + if metrics_output[2] is not None: evaluation_metrics.append(("validation/secondary_mbACC", metrics_output[2].mbACC())) diff --git a/yoeo/utils/dataclasses.py b/yoeo/utils/dataclasses.py index c2973cd..39422eb 100644 --- a/yoeo/utils/dataclasses.py +++ b/yoeo/utils/dataclasses.py @@ -10,10 +10,13 @@ class ClassNames: detection: List[str] segmentation: List[str] + base_footprint: List[str] @classmethod def load_from(cls, path: str) -> ClassNames: + print(path) file_content = cls._read_yaml_file(path) + print(file_content) class_names = cls._parse_yaml_file(file_content) return class_names @@ -21,7 +24,7 @@ def load_from(cls, path: str) -> ClassNames: @staticmethod def _parse_yaml_file(content: Dict[Any, Any]) -> ClassNames: return ClassNames(**content) - + @staticmethod def _read_yaml_file(path: str) -> Dict[Any, Any]: with open(path, "r") as f: diff --git a/yoeo/utils/datasets.py b/yoeo/utils/datasets.py index bc31f4a..774987a 100644 --- a/yoeo/utils/datasets.py +++ b/yoeo/utils/datasets.py @@ -111,7 +111,8 @@ def __getitem__(self, index): # Ignore warning if file is empty with warnings.catch_warnings(): warnings.simplefilter("ignore") - boxes = np.loadtxt(label_path).reshape(-1, 5) + boxes = np.loadtxt(label_path).reshape(-1, 7) # TODO use 5 if we do not use a base_footprint annotation + except Exception: print(f"Could not read label '{label_path}'.") return diff --git a/yoeo/utils/loss.py b/yoeo/utils/loss.py index f819b52..dce47f5 100644 --- a/yoeo/utils/loss.py +++ b/yoeo/utils/loss.py @@ -1,4 +1,5 @@ import math +from collections import defaultdict import torch import torch.nn as nn @@ -66,10 +67,14 @@ def compute_loss(combined_predictions, combined_targets, model): seg_loss = nn.CrossEntropyLoss()(seg_predictions[0], seg_targets).unsqueeze(0) # Add placeholder varables for the different losses - lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) + lcls = torch.zeros(1, device=device) + lbox = torch.zeros(1, device=device) + lobj = torch.zeros(1, device=device) + lbasepoint = torch.zeros(1, device=device) + lbasevisible = torch.zeros(1, device=device) # Build yolo targets - tcls, tbox, indices, anchors = build_targets(yolo_predictions, yolo_targets, model) # targets + yolo_targets = build_targets(yolo_predictions, yolo_targets, model) # targets # Define different loss functions classification BCEcls = nn.BCEWithLogitsLoss( @@ -77,10 +82,13 @@ def compute_loss(combined_predictions, combined_targets, model): BCEobj = nn.BCEWithLogitsLoss( pos_weight=torch.tensor([1.0], device=device)) + # Define the loss for regression + MSE = nn.MSELoss() + # Calculate losses for each yolo layer for layer_index, layer_predictions in enumerate(yolo_predictions): # Get image ids, anchors, grid index i and j for each target in the current yolo layer - b, anchor, grid_j, grid_i = indices[layer_index] + b, anchor, grid_j, grid_i = yolo_targets['indices'][layer_index] # Build empty object target tensor with the same shape as the object prediction tobj = torch.zeros_like(layer_predictions[..., 0], device=device) # target obj # Get the number of targets for this layer. @@ -96,11 +104,11 @@ def compute_loss(combined_predictions, combined_targets, model): # Apply sigmoid to xy offset predictions in each cell that has a target pxy = ps[:, :2].sigmoid() # Apply exponent to wh predictions and multiply with the anchor box that matched best with the label for each cell that has a target - pwh = torch.exp(ps[:, 2:4]) * anchors[layer_index] + pwh = torch.exp(ps[:, 2:4]) * yolo_targets['anchor'][layer_index] # Build box out of xy and wh pbox = torch.cat((pxy, pwh), 1) # Calculate CIoU or GIoU for each target with the predicted box for its cell + anchor - iou = bbox_iou(pbox.T, tbox[layer_index], x1y1x2y2=False, CIoU=True) + iou = bbox_iou(pbox.T, yolo_targets['box'][layer_index], x1y1x2y2=False, CIoU=True) # We want to minimize our loss so we and the best possible IoU is 1 so we take 1 - IoU and reduce it with a mean lbox += (1.0 - iou).mean() # iou loss @@ -113,10 +121,20 @@ def compute_loss(combined_predictions, combined_targets, model): if ps.size(1) - 5 > 1: # Hot one class encoding t = torch.zeros_like(ps[:, 5:], device=device) # targets - t[range(num_targets), tcls[layer_index]] = 1 + t[range(num_targets), yolo_targets['class'][layer_index]] = 1 # Use the tensor to calculate the BCE loss lcls += BCEcls(ps[:, 5:], t) # BCE + # Check if we need to calculate the base footprint loss + if 'base_footprint' in yolo_targets: + # Get the base footprints and the visibility of the base footprints + base_footprints, base_footprint_visible = yolo_targets['base_footprint'][layer_index] + # Calculate the MSE loss between the predicted base footprints and the target base footprints + if base_footprint_visible.any(): + lbasepoint += MSE(ps[:, -2:][base_footprint_visible], torch.arctan(base_footprints[base_footprint_visible])) + # Do binary classification of the visibility of the base footprints + lbasevisible += BCEcls(ps[:, -3], base_footprint_visible.type_as(ps)) + # Classification of the objectness the sequel # Calculate the BCE loss between the on the fly generated target and the network prediction lobj += BCEobj(layer_predictions[..., 4], tobj) # obj loss @@ -125,68 +143,100 @@ def compute_loss(combined_predictions, combined_targets, model): lbox *= 0.2 lobj *= 10.0 lcls *= 0.05 + lbasepoint *= 0.5 + lbasevisible *= 0.1 # Merge losses - loss = lbox + lobj + lcls + seg_loss + loss = lbox + lobj + lcls + seg_loss + lbasepoint + lbasevisible - return loss, to_cpu(torch.cat((lbox, lobj, lcls, seg_loss, loss))) + return loss, to_cpu(torch.cat((lbox, lobj, lcls, lbasepoint, lbasevisible, seg_loss, loss))) def build_targets(p, targets, model): # Build targets for compute_loss(), input targets(image,class,x,y,w,h) na, nt = len(model.yolo_layers[0].anchors), targets.shape[0] # number of anchors, targets - tcls, tbox, indices, anch = [], [], [], [] - gain = torch.ones(7, device=targets.device) # normalized to gridspace gain + + # Check if we need to build targets for base footprints + base_footprint_present = targets.shape[1] > 5 + + # Initialize the output dictionary + output = defaultdict(list) + # Make a tensor that iterates 0-2 for 3 anchors and repeat that as many times as we have target boxes - ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) + anchor_index = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # Copy target boxes anchor size times and append an anchor index to each copy the anchor index is also expressed by the new first dimension - targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) + targets = torch.cat((targets.repeat(na, 1, 1), anchor_index[:, :, None]), 2) for i, yolo_layer in enumerate(model.yolo_layers): # Scale anchors by the yolo grid cell size so that an anchor with the size of the cell would result in 1 anchors = yolo_layer.anchors / yolo_layer.stride + + # Get the grid size of the yolo layer + grid_size_x, grid_size_y = p[i].shape[3], p[i].shape[2] + # Add the number of yolo cells in this layer the gain tensor - # The gain tensor matches the collums of our targets (img id, class, x, y, w, h, anchor id) - gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain + # The gain tensor matches the columns of our targets (img id, class, x, y, w, h, [base footprint x, base footprint y,] anchor id) + if base_footprint_present: + # Keep non spatial values the same + gain = [1, 1, grid_size_x, grid_size_y, grid_size_x, grid_size_y, grid_size_x, grid_size_y, 1] + else: + # Keep non spatial values the same + gain = [1, 1, grid_size_x, grid_size_y, grid_size_x, grid_size_y, 1] + # Create a tensor out of the gain list + gain = torch.tensor(gain, device=targets.device, dtype=torch.long) # Scale targets by the number of yolo layer cells, they are now in the yolo cell coordinate system t = targets * gain - # Check if we have targets - if nt: - t = t.squeeze(0) - # Check if we use anchor boxes - if na != 1: - # Calculate ratio between anchor and target box for both width and height - r = t[:, :, 4:6] / anchors[:, None] - # Select the ratios that have the highest divergence in any axis and check if the ratio is less than 4 - j = torch.max(r, 1. / r).max(2)[0] < 4 - # Only use targets that have the correct ratios for their anchors - # That means we only keep ones that have a matching anchor and we loose the anchor dimension - # The anchor id is still saved in the 7th value of each target - t = t[j] - else: - t = targets[0] + + # Remove anchor dimension if possible + t = t.squeeze(0) + # Check if we use anchor boxes + if na != 1: + # Calculate ratio between anchor and target box for both width and height + r = t[:, :, 4:6] / anchors[:, None] + # Select the ratios that have the highest divergence in any axis and check if the ratio is less than 4 + j = torch.max(r, 1. / r).max(2)[0] < 4 + # Only use targets that have the correct ratios for their anchors + # That means we only keep ones that have a matching anchor and we loose the anchor dimension + # The anchor id is still saved in the 7th value of each target + t = t[j] # Extract image id in batch and class id b, c = t[:, :2].long().T # We isolate the target cell associations. - # x, y, w, h are allready in the cell coordinate system meaning an x = 1.2 would be 1.2 times cellwidth - gxy = t[:, 2:4] - gwh = t[:, 4:6] # grid wh + # x, y, w, h are already in the cell coordinate system meaning an x = 1.2 would be 1.2 times cell width + gxy = t[:, 2:4] # xy in grid cell coordinates + gwh = t[:, 4:6] # wh in grid cell coordinates # Cast to int to get an cell index e.g. 1.2 gets associated to cell 1 gij = gxy.long() # Isolate x and y index dimensions gi, gj = gij.T # grid xy indices + # Get the x and y coordinates in the coordinate system of the cell + rxy = gxy - gij # Convert anchor indexes to int - a = t[:, 6].long() + a = t[:, -1].long() + + # Check if we need to build targets for base footprints + if base_footprint_present: + # Extract the base footprints from the target tensor + gbxy = t[:, 6:8] + # Get the base footprints in the coordinate system of the cell of the corresponding object + rbxy = gbxy - gij + # Check if we have any nan values in the base footprints + # We use them to indicate that the base footprint is not visible + # We want to ignore those in the loss calculation and have them as a separate classification target + base_footprint_visible = ~torch.isnan(rbxy).any(1) + # Add the base footprints to the target lists + output['base_footprint'].append((rbxy, base_footprint_visible)) + # Add target tensors for this yolo layer to the output lists # Add to index list and limit index range to prevent out of bounds - indices.append((b, a, gj.clamp_(0, gain[3].long() - 1), gi.clamp_(0, gain[2].long() - 1))) + output['indices'].append((b, a, gj.clamp_(0, grid_size_x - 1), gi.clamp_(0, grid_size_y - 1))) # Add to target box list and convert box coordinates from global grid coordinates to local offsets in the grid cell - tbox.append(torch.cat((gxy - gij, gwh), 1)) # box + output['box'].append(torch.cat((rxy, gwh), 1)) # box # Add correct anchor for each target to the list - anch.append(anchors[a]) + output['anchor'].append(anchors[a]) # Add class for each target to the list - tcls.append(c) + output['class'].append(c) - return tcls, tbox, indices, anch + return output diff --git a/yoeo/utils/transforms.py b/yoeo/utils/transforms.py index 800c457..e9f87a7 100644 --- a/yoeo/utils/transforms.py +++ b/yoeo/utils/transforms.py @@ -5,6 +5,7 @@ import imgaug.augmenters as iaa from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage from imgaug.augmentables.segmaps import SegmentationMapsOnImage +from imgaug.augmentables.kps import Keypoint, KeypointsOnImage from .utils import xywh2xyxy_np import torchvision.transforms as transforms @@ -18,29 +19,52 @@ def __call__(self, data): # Unpack data img, boxes, seg = data - # Convert xywh to xyxy + # Extract base_footprints if they exist boxes = np.array(boxes) - boxes[:, 1:] = xywh2xyxy_np(boxes[:, 1:]) + base_footprints = boxes[:, 5:] + base_footprints_present = base_footprints.shape[1] > 0 + + # Replace -1 with nan as our placeholder for no base_footprint + base_footprints[base_footprints==-1] = np.nan + + # Convert xywh to xyxy + boxes[:, 1:5] = xywh2xyxy_np(boxes[:, 1:5]) # Convert bounding boxes to imgaug bounding_boxes = BoundingBoxesOnImage( - [BoundingBox(*box[1:], label=box[0]) for box in boxes], + [BoundingBox(*box[1:5], label=box[0]) for box in boxes], shape=img.shape) + # Convert base_footprints to imgaug + kps = KeypointsOnImage([ + Keypoint(*base_footprint) for base_footprint in base_footprints + ] if base_footprints_present else [], + shape=img.shape + ) + # Convert sementations to imgaug segmentation_mask = SegmentationMapsOnImage(seg, shape=img.shape) # Apply augmentations - img, bounding_boxes, segmentation_mask = self.augmentations( + img, bounding_boxes, segmentation_mask, kps = self.augmentations( image=img, bounding_boxes=bounding_boxes, - segmentation_maps=segmentation_mask) + segmentation_maps=segmentation_mask, + keypoints=kps) # Clip out of image boxes bounding_boxes = bounding_boxes.clip_out_of_image() + # Disable out of image keypoints + for kp in kps.keypoints: + if kp.is_out_of_image(img): + #print(f"Keypoint out of image: {kp.x}, {kp.y} for image of shape {img.shape}") + kp.x = np.nan + kp.y = np.nan + # Convert bounding boxes back to numpy - boxes = np.zeros((len(bounding_boxes), 5)) + box_dimension = 5 if not base_footprints_present else 7 + boxes = np.zeros((len(bounding_boxes), box_dimension)) for box_idx, box in enumerate(bounding_boxes): # Extract coordinates for unpadded + unscaled image x1 = box.x1 @@ -55,6 +79,11 @@ def __call__(self, data): boxes[box_idx, 3] = (x2 - x1) boxes[box_idx, 4] = (y2 - y1) + # Add base_footprint if it exists + if base_footprints_present: + boxes[box_idx, 5] = kps.keypoints[box_idx].x + boxes[box_idx, 6] = kps.keypoints[box_idx].y + # Convert segmentation back to numpy seg = segmentation_mask.get_arr() @@ -68,8 +97,13 @@ def __init__(self, ): def __call__(self, data): img, boxes, seg = data h, w, _ = img.shape + # Transform boxes to relative coordinates boxes[:,[1, 3]] /= w boxes[:,[2, 4]] /= h + # Transform the base_footprint of the box if it exists + if boxes.shape[1] > 5: + boxes[:, 5] /= w + boxes[:, 6] /= h return img, boxes, seg @@ -80,8 +114,13 @@ def __init__(self, ): def __call__(self, data): img, boxes, seg = data h, w, _ = img.shape + # Transform boxes to absolute coordinates boxes[:,[1, 3]] *= w boxes[:,[2, 4]] *= h + # Transform the base_footprint of the box if it exists + if boxes.shape[1] > 5: + boxes[:, 5] *= w + boxes[:, 6] *= h return img, boxes, seg @@ -102,9 +141,9 @@ def __call__(self, data): img, boxes, seg = data # Extract image as PyTorch tensor img = transforms.ToTensor()(img) - seg = transforms.ToTensor()(seg) * 255 # Because troch maps this to 0-1 instead of 0-255 + seg = transforms.ToTensor()(seg) * 255 # Because torch maps this to 0-1 instead of 0-255 - bb_targets = torch.zeros((len(boxes), 6)) + bb_targets = torch.zeros((len(boxes), 1 + boxes.shape[1])) bb_targets[:, 1:] = transforms.ToTensor()(boxes) return img, bb_targets, seg diff --git a/yoeo/utils/utils.py b/yoeo/utils/utils.py index de62950..843a6be 100644 --- a/yoeo/utils/utils.py +++ b/yoeo/utils/utils.py @@ -93,6 +93,7 @@ def rescale_boxes_to_original_padded_img_size(boxes, output_img_size: int, origi scale_factor = original_max_size / output_img_size boxes[:, 0:4] = boxes[:, 0:4] * scale_factor + boxes[:, -2:] = boxes[:, -2:] * scale_factor return boxes @@ -117,6 +118,8 @@ def unpad_box_coordinates(boxes, original_img_size: Tuple[int, int]): boxes[:, 1] = boxes[:, 1] - padding_top boxes[:, 2] = boxes[:, 2] - padding_left boxes[:, 3] = boxes[:, 3] - padding_top + boxes[:, -2] = boxes[:, -2] - padding_left + boxes[:, -1] = boxes[:, -1] - padding_top return boxes @@ -164,7 +167,7 @@ def rescale_segmentation_to_original_padded_img_size(segmentation, original_max_ def unpad_segmentation(segmentation, original_img_size: Tuple[int, int]): """ Remove paddings, thereby restoring the original image size - + :param segmentation: segmentation output :type segmentation: torch.Tensor with shape (1, height, width) and height == width :param original_img_size: original image size (height, width) @@ -288,9 +291,9 @@ def compute_ap(recall, precision): return ap -def get_batch_statistics(outputs, - targets, - iou_threshold, +def get_batch_statistics(outputs, + targets, + iou_threshold, group_config: Optional[GroupConfig] = None ) -> Tuple[List, Optional[Metric]]: """ @@ -321,7 +324,9 @@ class ('GroupConfig.surrogate_id') for batch statistics evaluation and evaluated output = outputs[sample_i] pred_boxes = output[:, :4] pred_scores = output[:, 4] - pred_labels = output[:, -1] + pred_labels = output[:, 5] + pred_base_footprint_visibility = output[:, -3] + pred_base_footprint_point = output[:, -2:] if grouping_active: sec_pred_labels = compute_secondary_labels(pred_labels, group_ids) @@ -329,6 +334,15 @@ class ('GroupConfig.surrogate_id') for batch statistics evaluation and evaluated true_positives = np.zeros(pred_boxes.shape[0]) + # Base footprint visibility confusion matrix + base_bootprint_visibility_tp = 0 + base_bootprint_visibility_fn = 0 + base_bootprint_visibility_fp = 0 + base_bootprint_visibility_tn = 0 + + # Base footprint point distances + base_footprint_point_distances = [] + annotations = targets[targets[:, 0] == sample_i][:, 1:] target_labels = annotations[:, 0] if len(annotations) else [] @@ -338,7 +352,7 @@ class ('GroupConfig.surrogate_id') for batch statistics evaluation and evaluated if len(annotations): detected_boxes = [] - target_boxes = annotations[:, 1:] + target_boxes = annotations[:, 1:-2] for pred_i, (pred_box, pred_label) in enumerate(zip(pred_boxes, pred_labels)): # If targets are found break @@ -364,17 +378,47 @@ class ('GroupConfig.surrogate_id') for batch statistics evaluation and evaluated true_positives[pred_i] = 1 detected_boxes += [box_index] + # Check if the base footprint visibility is correct + target_box_base_footprint_visibility = ~torch.isnan(annotations[box_index, -2:]).any() + pred_box_base_footprint_visibility = pred_base_footprint_visibility[pred_i] > 0.5 # TODO configurable threshold + + if target_box_base_footprint_visibility and pred_box_base_footprint_visibility: + base_bootprint_visibility_tp += 1 + + # If the base footprint visibility is correct we can look at the x and y coordinates + target_base_footprint_point = annotations[box_index, -2:] + + # Calculate the euclidean distance between the target and predicted base footprint point + base_footprint_point_distance = torch.dist(target_base_footprint_point, pred_base_footprint_point[pred_i]) + + # Base footprint point distances + base_footprint_point_distances.append(base_footprint_point_distance) + elif target_box_base_footprint_visibility and not pred_box_base_footprint_visibility: + base_bootprint_visibility_fn += 1 + elif not target_box_base_footprint_visibility and pred_box_base_footprint_visibility: + base_bootprint_visibility_fp += 1 + else: + base_bootprint_visibility_tn += 1 + if grouping_active: sec_pred_label = sec_pred_labels[pred_i] if pred_label in group_ids: secondary_metric.update(sec_pred_label.int(), sec_target_labels[box_index].int()) - batch_metrics.append([true_positives, pred_scores, pred_labels]) + batch_metrics.append([ + true_positives, + pred_scores, + pred_labels, + np.array([base_bootprint_visibility_tp]), + np.array([base_bootprint_visibility_fn]), + np.array([base_bootprint_visibility_fp]), + np.array([base_bootprint_visibility_tn]), + np.array(base_footprint_point_distances)]) return batch_metrics, secondary_metric -def compute_secondary_labels(labels: torch.tensor, group_ids: torch.tensor) -> torch.tensor: +def compute_secondary_labels(labels: torch.tensor, group_ids: torch.tensor) -> torch.tensor: secondary_labels = labels.clone() # We replace the actual class labels with values from {0, ...} for classes that should be grouped into a @@ -382,7 +426,7 @@ def compute_secondary_labels(labels: torch.tensor, group_ids: torch.tensor) -> t for idx, squeeze_id in enumerate(group_ids): # Replace label with value in {0, ...} secondary_labels[labels == squeeze_id] = idx - + # Replace all other labels with -1 secondary_labels[torch.logical_not(torch.isin(labels, group_ids))] = -1 @@ -469,14 +513,14 @@ def box_area(box): def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, group_config: Optional[GroupConfig] = None): """ - Performs Non-Maximum Suppression (NMS) on inference results. If 'group_config' is not 'None', the contained + Performs Non-Maximum Suppression (NMS) on inference results. If 'group_config' is not 'None', the contained classes will be treated as one class ('GroupConfig.surrogate_id') during non-maximum supression. - + Returns: detections with shape: nx6 (x1, y1, x2, y2, conf, cls) """ - nc = prediction.shape[2] - 5 # number of classes + nc = prediction.shape[2] - 5 - 3 # number of classes # Settings # (pixels) minimum and maximum box width and height @@ -487,7 +531,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) t = time.time() - output = [torch.zeros((0, 6), device="cpu")] * prediction.shape[0] + output = [torch.zeros((0, 9), device="cpu")] * prediction.shape[0] if group_config: group_ids = torch.tensor(group_config.group_ids, device=prediction.device, dtype=prediction.dtype) @@ -501,18 +545,18 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non continue # Compute conf - x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf + x[:, 5:-3] *= x[:, 4:5] # conf = obj_conf * cls_conf # Box (center x, center y, width, height) to (x1, y1, x2, y2) box = xywh2xyxy(x[:, :4]) - # Detections matrix nx6 (xyxy, conf, cls) + # Detections matrix nx6 (xyxy, conf, cls, base_footprint_visiblity, base_footprint_x, base_footprint_y) if multi_label: - i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T - x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) + i, j = (x[:, 5:-3] > conf_thres).nonzero(as_tuple=False).T + x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float(), x[i, -3:]), 1) else: # best class only - conf, j = x[:, 5:].max(1, keepdim=True) - x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] + conf, j = x[:, 5:-3].max(1, keepdim=True) + x = torch.cat((box, conf, j.float(), x[i, -3:]), 1)[conf.view(-1) > conf_thres] # Filter by class if classes is not None: @@ -530,7 +574,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non if group_config is None: c = x[:, 5:6] * max_wh # classes else: - # If for example multiple robot classes are present, all robot classes are treated as one class in order + # If for example multiple robot classes are present, all robot classes are treated as one class in order # to perform nms across all classes and not per class. For this, all robot classes get the same offset. c = torch.clone(x[:, 5:6]) c[torch.isin(c, group_ids)] = group_config.surrogate_id