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