From 5f3455482176ca84795f61d2248f1dc031664db0 Mon Sep 17 00:00:00 2001 From: Simon Date: Thu, 18 Jan 2024 19:55:33 +0100 Subject: [PATCH 1/3] add Negative Data training and testing --- yoeo/test.py | 33 +++++++++++++++++++++++++++------ yoeo/train.py | 31 ++++++++++++++++++++++++++----- yoeo/utils/datasets.py | 39 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 91 insertions(+), 12 deletions(-) diff --git a/yoeo/test.py b/yoeo/test.py index d266572..b0d1416 100755 --- a/yoeo/test.py +++ b/yoeo/test.py @@ -10,18 +10,18 @@ from terminaltables import AsciiTable import torch -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, ConcatDataset from torch.autograd import Variable from yoeo.models import load_model from yoeo.utils.utils import load_classes, ap_per_class, get_batch_statistics, non_max_suppression, to_cpu, xywh2xyxy, \ print_environment_info, seg_iou -from yoeo.utils.datasets import ListDataset +from yoeo.utils.datasets import ListDataset, NegativeDataset from yoeo.utils.transforms import DEFAULT_TRANSFORMS from yoeo.utils.parse_config import parse_data_config -def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_size=8, img_size=416, +def evaluate_model_file(model_path, weights_path, img_path, class_names, negative_img_dir="", negative_data_fraction=0.0, batch_size=8, img_size=416, n_cpu=8, iou_thres=0.5, conf_thres=0.5, nms_thres=0.5, verbose=True, robot_class_ids: Optional[List[int]] = None): """Evaluate model on validation dataset. @@ -34,6 +34,10 @@ def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_s :type img_path: str :param class_names: Dict containing detection and segmentation class names :type class_names: Dict + :param negative_img_dir: Path to negative image folder, defaults to "" + :type negative_img_dir: str + :param negative_data_fraction: Fraction of negative data relative to positive data, defaults to 0.0 + :type negative_data_fraction: float :param batch_size: Size of each image batch, defaults to 8 :type batch_size: int, optional :param img_size: Size of each image dimension for yolo, defaults to 416 @@ -53,7 +57,7 @@ def evaluate_model_file(model_path, weights_path, img_path, class_names, batch_s :return: Returns precision, recall, AP, f1, ap_class """ dataloader = _create_validation_data_loader( - img_path, batch_size, img_size, n_cpu) + img_path, negative_img_dir, negative_data_fraction, batch_size, img_size, n_cpu) model = load_model(model_path, weights_path) metrics_output, seg_class_ious = _evaluate( model, @@ -180,12 +184,16 @@ def seg_iou_mean_without_nan(seg_iou: List[float]) -> np.ndarray: return yolo_metrics_output, seg_class_ious -def _create_validation_data_loader(img_path, batch_size, img_size, n_cpu): +def _create_validation_data_loader(img_path, negative_img_dir, negative_data_fraction, batch_size, img_size, n_cpu): """ Creates a DataLoader for validation. :param img_path: Path to file containing all paths to validation images. :type img_path: str + :param negative_img_dir: Path to negative image folder + :type negative_img_dir: str + :param negative_data_fraction: Fraction of negative data relative to positive data + :type negative_data_fraction: float :param batch_size: Size of each image batch :type batch_size: int :param img_size: Size of each image dimension for yolo @@ -196,8 +204,19 @@ def _create_validation_data_loader(img_path, batch_size, img_size, n_cpu): :rtype: DataLoader """ dataset = ListDataset(img_path, img_size=img_size, multiscale=False, transform=DEFAULT_TRANSFORMS) + + dataset_len = len(dataset) + negative_dataset_len = int(negative_data_fraction*dataset_len) + + negative_dataset = NegativeDataset( + negative_img_dir, + img_size=img_size, + negative_dataset_max_len=negative_dataset_len) + + concat_dataset = ConcatDataset([dataset, negative_dataset]) + dataloader = DataLoader( - dataset, + concat_dataset, batch_size=batch_size, shuffle=False, num_workers=n_cpu, @@ -214,6 +233,8 @@ def run(): parser.add_argument("-w", "--weights", type=str, default="weights/yoeo.pth", help="Path to weights or checkpoint file (.weights or .pth)") parser.add_argument("-d", "--data", type=str, default="config/torso.data", help="Path to data config file (.data)") + parser.add_argument("-n", "--negative_data_dir", default='', type=str, help="Path to negative data directory") + parser.add_argument("--negative_data_fraction", default=0, type=float, help="Fraction of negative data relative to positive data (default=0.0)") parser.add_argument("-b", "--batch_size", type=int, default=8, help="Size of each image batch") parser.add_argument("-v", "--verbose", action='store_true', help="Makes the validation more verbose") parser.add_argument("--img_size", type=int, default=416, help="Size of each image dimension for yolo") diff --git a/yoeo/train.py b/yoeo/train.py index ef328cd..5794fdb 100755 --- a/yoeo/train.py +++ b/yoeo/train.py @@ -9,7 +9,7 @@ import numpy as np import torch -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, ConcatDataset import torch.optim as optim from torch.autograd import Variable @@ -18,7 +18,7 @@ from yoeo.models import load_model from yoeo.utils.logger import Logger from yoeo.utils.utils import to_cpu, load_classes, print_environment_info, provide_determinism, worker_seed_set -from yoeo.utils.datasets import ListDataset +from yoeo.utils.datasets import ListDataset, NegativeDataset from yoeo.utils.augmentations import AUGMENTATION_TRANSFORMS from yoeo.utils.transforms import DEFAULT_TRANSFORMS from yoeo.utils.parse_config import parse_data_config @@ -30,11 +30,15 @@ from torchsummary import summary -def _create_data_loader(img_path, batch_size, img_size, n_cpu, multiscale_training=False): +def _create_data_loader(img_path, negative_img_dir, negative_data_fraction, batch_size, img_size, n_cpu, multiscale_training=False): """Creates a DataLoader for training. :param img_path: Path to file containing all paths to training images. :type img_path: str + :param negative_img_dir: Path to negative image folder + :type negative_img_dir: str + :param negative_data_fraction: Fraction of negative data relative to positive data + :type negative_data_fraction: float :param batch_size: Size of each image batch :type batch_size: int :param img_size: Size of each image dimension for yolo @@ -51,8 +55,20 @@ def _create_data_loader(img_path, batch_size, img_size, n_cpu, multiscale_traini img_size=img_size, multiscale=multiscale_training, transform=AUGMENTATION_TRANSFORMS) + + dataset_len = len(dataset) + negative_dataset_len = int(negative_data_fraction*dataset_len) + + negative_dataset = NegativeDataset( + negative_img_dir, + img_size=img_size, + transform=AUGMENTATION_TRANSFORMS, + negative_dataset_max_len=negative_dataset_len) + + concat_dataset = ConcatDataset([dataset, negative_dataset]) + dataloader = DataLoader( - dataset, + concat_dataset, batch_size=batch_size, shuffle=True, num_workers=n_cpu, @@ -61,12 +77,13 @@ def _create_data_loader(img_path, batch_size, img_size, n_cpu, multiscale_traini worker_init_fn=worker_seed_set) return dataloader - def run(): print_environment_info() parser = argparse.ArgumentParser(description="Trains the YOLO model.") parser.add_argument("-m", "--model", type=str, default="config/yoeo.cfg", help="Path to model definition file (.cfg)") parser.add_argument("-d", "--data", type=str, default="config/torso.data", help="Path to data config file (.data)") + parser.add_argument("-n", "--negative_data_dir", default='', type=str, help="Path to negative data directory") + parser.add_argument("--negative_data_fraction", default=0, type=float, help="Fraction of negative data relative to positive data (default=0.0)") parser.add_argument("-e", "--epochs", type=int, default=300, help="Number of epochs") parser.add_argument("-v", "--verbose", action='store_true', help="Makes the training more verbose") parser.add_argument("--n_cpu", type=int, default=8, help="Number of cpu threads to use during batch generation") @@ -128,6 +145,8 @@ def run(): # Load training dataloader dataloader = _create_data_loader( train_path, + args.negative_data_dir, + args.negative_data_fraction, mini_batch_size, model.hyperparams['height'], args.n_cpu, @@ -136,6 +155,8 @@ def run(): # Load validation dataloader validation_dataloader = _create_validation_data_loader( valid_path, + args.negative_data_dir, + args.negative_data_fraction, mini_batch_size, model.hyperparams['height'], args.n_cpu) diff --git a/yoeo/utils/datasets.py b/yoeo/utils/datasets.py index bc31f4a..43ebf18 100644 --- a/yoeo/utils/datasets.py +++ b/yoeo/utils/datasets.py @@ -54,6 +54,44 @@ def __getitem__(self, index): def __len__(self): return len(self.files) + + +class NegativeDataset(Dataset): + def __init__(self, folder_path, img_size=416, transform=None,negative_dataset_max_len=0): + self.img_size = img_size + self.transform = transform + self.negative_dataset_max_len = negative_dataset_max_len + if folder_path: + self.files = sorted(glob.glob("%s/*.*" % folder_path))[:self.negative_dataset_max_len] + else: + self.files = [] + + def __getitem__(self, index): + img_path = self.files[index % len(self.files)] + img = np.array( + Image.open(img_path).convert('RGB'), + dtype=np.uint8) + + # Label Placeholder + bb_targets = np.zeros((1, 5)) + mask_targets = np.zeros_like(img) + + # ----------- + # Transform + # ----------- + if self.transform: + try: + img, bb_targets, mask_targets = self.transform( + (img, bb_targets, mask_targets) + ) + except Exception as e: + print(f"Could not apply transform.") + raise e + + return img_path, img, bb_targets, mask_targets + + def __len__(self): + return len(self.files) class ListDataset(Dataset): @@ -138,7 +176,6 @@ def __getitem__(self, index): except Exception as e: print(f"Could not apply transform.") raise e - return return img_path, img, bb_targets, mask_targets From 634356f136c72a7c708770018f6204e70336ecb3 Mon Sep 17 00:00:00 2001 From: Simon Date: Thu, 1 Feb 2024 19:14:34 +0100 Subject: [PATCH 2/3] Fix validation dataset + fix unbound variables --- yoeo/test.py | 2 +- yoeo/utils/datasets.py | 13 +++++-------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/yoeo/test.py b/yoeo/test.py index b0d1416..e5157ba 100755 --- a/yoeo/test.py +++ b/yoeo/test.py @@ -211,10 +211,10 @@ def _create_validation_data_loader(img_path, negative_img_dir, negative_data_fra negative_dataset = NegativeDataset( negative_img_dir, img_size=img_size, + transform=DEFAULT_TRANSFORMS, negative_dataset_max_len=negative_dataset_len) concat_dataset = ConcatDataset([dataset, negative_dataset]) - dataloader = DataLoader( concat_dataset, batch_size=batch_size, diff --git a/yoeo/utils/datasets.py b/yoeo/utils/datasets.py index 43ebf18..710e840 100644 --- a/yoeo/utils/datasets.py +++ b/yoeo/utils/datasets.py @@ -95,7 +95,7 @@ def __len__(self): class ListDataset(Dataset): - def __init__(self, list_path, img_size=416, multiscale=True, transform=None): + def __init__(self, list_path, img_size: int =416, multiscale=True, transform=None): with open(list_path, "r") as file: self.img_files = file.readlines() @@ -119,7 +119,7 @@ def __init__(self, list_path, img_size=416, multiscale=True, transform=None): mask_file = os.path.splitext(mask_file)[0] + '.png' self.mask_files.append(mask_file) - self.img_size = img_size + self.img_size: int = img_size self.max_objects = 100 self.multiscale = multiscale self.min_size = self.img_size - 3 * 32 @@ -132,9 +132,8 @@ def __getitem__(self, index): # --------- # Image # --------- + img_path = self.img_files[index % len(self.img_files)].rstrip() try: - img_path = self.img_files[index % len(self.img_files)].rstrip() - img = np.array(Image.open(img_path).convert('RGB'), dtype=np.uint8) except Exception: print(f"Could not read image '{img_path}'.") @@ -143,9 +142,8 @@ def __getitem__(self, index): # --------- # Label # --------- + label_path = self.label_files[index % len(self.img_files)].rstrip() try: - label_path = self.label_files[index % len(self.img_files)].rstrip() - # Ignore warning if file is empty with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -157,8 +155,8 @@ def __getitem__(self, index): # --------- # Segmentation Mask # --------- + mask_path = self.mask_files[index % len(self.img_files)].rstrip() try: - mask_path = self.mask_files[index % len(self.img_files)].rstrip() # Load segmentation mask as numpy array mask = np.array(Image.open(mask_path).convert('RGB')) except FileNotFoundError as e: @@ -184,7 +182,6 @@ def collate_fn(self, batch): # Drop invalid images batch = [data for data in batch if data is not None] - paths, imgs, bb_targets, mask_targets = list(zip(*batch)) # Selects new image size every tenth batch From 1421108d327cec877b2e05f5caf8c10e9745cc2b Mon Sep 17 00:00:00 2001 From: Simon Date: Thu, 1 Feb 2024 19:38:32 +0100 Subject: [PATCH 3/3] fix unnecessary load_classes import --- yoeo/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yoeo/train.py b/yoeo/train.py index 3061fcb..f9c7aa3 100755 --- a/yoeo/train.py +++ b/yoeo/train.py @@ -15,7 +15,7 @@ from yoeo.models import load_model from yoeo.utils.logger import Logger -from yoeo.utils.utils import to_cpu, load_classes, print_environment_info, provide_determinism, worker_seed_set +from yoeo.utils.utils import to_cpu, print_environment_info, provide_determinism, worker_seed_set from yoeo.utils.datasets import ListDataset, NegativeDataset from yoeo.utils.dataclasses import ClassNames from yoeo.utils.class_config import ClassConfig