diff --git a/datasets.py b/datasets.py index 2da3cb80..88eb1a80 100644 --- a/datasets.py +++ b/datasets.py @@ -28,7 +28,8 @@ def __init__( use_train_aug=False, train=False, no_mosaic=False, - square_training=False + square_training=False, + discard_negative_example = True ): self.transforms = transforms self.use_train_aug = use_train_aug @@ -41,59 +42,42 @@ def __init__( self.square_training = square_training self.mosaic_border = [-img_size // 2, -img_size // 2] self.image_file_types = ['*.jpg', '*.jpeg', '*.png', '*.ppm', '*.JPG'] - self.all_image_paths = [] - - # get all the image paths in sorted order - for file_type in self.image_file_types: - self.all_image_paths.extend(glob.glob(os.path.join(self.images_path, file_type))) + self.all_images = [] + self.all_annot_paths = glob.glob(os.path.join(self.labels_path, '*.xml')) - self.all_images = [image_path.split(os.path.sep)[-1] for image_path in self.all_image_paths] - self.all_images = sorted(self.all_images) - # Remove all annotations and images when no object is present. - self.read_and_clean() - def read_and_clean(self): + self.read_and_clean(discard_negative_example) + + def read_and_clean(self, discard_negative_example): # Discard any images and labels when the XML - # file does not contain any object. - for annot_path in self.all_annot_paths: - tree = et.parse(annot_path) + # file does not contain any object, if negative example are not given + def check_path_and_save_image(path): + tree = et.parse(path) root = tree.getroot() - object_present = False - for member in root.findall('object'): - if member.find('bndbox'): - object_present = True - if object_present == False: - image_name = annot_path.split(os.path.sep)[-1].split('.xml')[0] - image_root = self.all_image_paths[0].split(os.path.sep)[:-1] - # remove_image = f"{'/'.join(image_root)}/{image_name}.jpg" - - # TODO Is this code necessary? - # for img_type in self.image_file_types: - # remove_image = os.path.join(os.sep.join(image_root), image_name+img_type.replace("*","")) - # if remove_image in self.all_image_paths: - # print(f"Removing {annot_path} and corresponding {remove_image}") - # self.all_annot_paths.remove(annot_path) - # self.all_image_paths.remove(remove_image) - # break - - # Discard any image file when no annotation file - # is not found for the image. - for image_name in self.all_images: - possible_xml_name = os.path.join(self.labels_path, os.path.splitext(image_name)[0]+'.xml') - if possible_xml_name not in self.all_annot_paths: - print(f"{possible_xml_name} not found...") - print(f"Removing {image_name} image") - # items = [item for item in items if item != element] - self.all_images = [image_instance for image_instance in self.all_images if image_instance != image_name] - # self.all_images.remove(image_name) - - # for image_path in self.all_image_paths: - # image_name = image_path.split(os.path.sep)[-1].split('.jpg')[0] - # possible_xml_name = f"{self.labels_path}/{image_name.split('.jpg')[0]}.xml" - # if possible_xml_name not in self.all_annot_paths: - # print(f"{possible_xml_name} not found...") - # print(f"Removing {image_name} image") - # self.all_image_paths.remove(image_path) + image_name = root.findtext("filename") + image_path = os.path.join(self.images_path, image_name) + discard_path = False + if not os.path.exists(image_path): + print(f"Image {image_path} associated to {path} not found...") + print(f"Discarding {path}...") + discard_path = True + + if not discard_path and discard_negative_example: + object_present = False + for member in root.findall('object'): + if member.find('bndbox'): + object_present = True + break + if not object_present: + print(f"File {path} contains no object. Discarding xml file and image...") + discard_path = True + #If i don't discard xml than save the image + if not discard_path: + self.all_images.append(image_name) + return discard_path + + self.all_annot_paths = list(filter(check_path_and_save_image, self.all_annot_paths )) + def resize(self, im, square=False): if square: @@ -159,7 +143,7 @@ def load_image_and_labels(self, index): ymax_final = (ymax/image_height)*image_resized.shape[0] boxes.append([xmin_final, ymin_final, xmax_final, ymax_final]) - + # Bounding box to tensor. boxes_length = len(boxes) boxes = torch.as_tensor(boxes, dtype=torch.float32) @@ -320,7 +304,8 @@ def create_train_dataset( classes, use_train_aug=False, no_mosaic=False, - square_training=False + square_training=False, + discard_negative=True ): train_dataset = CustomDataset( train_dir_images, @@ -331,7 +316,8 @@ def create_train_dataset( use_train_aug=use_train_aug, train=True, no_mosaic=no_mosaic, - square_training=square_training + square_training=square_training, + discard_negative_example=discard_negative ) return train_dataset def create_valid_dataset( @@ -339,7 +325,8 @@ def create_valid_dataset( valid_dir_labels, img_size, classes, - square_training=False + square_training=False, + discard_negative=True ): valid_dataset = CustomDataset( valid_dir_images, @@ -349,7 +336,8 @@ def create_valid_dataset( get_valid_transform(), train=False, no_mosaic=True, - square_training=square_training + square_training=square_training, + discard_negative_example=discard_negative ) return valid_dataset diff --git a/eval.py b/eval.py index eb3340d1..cb902508 100644 --- a/eval.py +++ b/eval.py @@ -55,6 +55,12 @@ default=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'), help='computation/training device, default is GPU if GPU present' ) + parser.add_argument( + '-dn', '--discard-negative', + dest='discard_negative', + action='store_true', + help='pass this if you want to discard images with no objects' + ) parser.add_argument( '-v', '--verbose', action='store_true', @@ -105,7 +111,9 @@ VALID_DIR_LABELS, IMAGE_SIZE, COCO_91_CLASSES, - square_training=args['square_training'] + square_training=args['square_training'], + discard_negative=args["discard_negative"] + ) # Load weights. @@ -118,7 +126,8 @@ VALID_DIR_LABELS, IMAGE_SIZE, CLASSES, - square_training=args['square_training'] + square_training=args['square_training'], + discard_negative=args["discard_negative"] ) model.to(DEVICE).eval() diff --git a/train.py b/train.py index bfed53b9..d19f0c02 100644 --- a/train.py +++ b/train.py @@ -116,6 +116,12 @@ def parse_opt(): action='store_true', help='pass this to not to use mosaic augmentation' ) + parser.add_argument( + '-dn', '--discard-negative', + dest='discard_negative', + action='store_true', + help='pass this if you want to discard images with no objects' + ) parser.add_argument( '-uta', '--use-train-aug', dest='use_train_aug', @@ -227,14 +233,16 @@ def main(args): CLASSES, use_train_aug=args['use_train_aug'], no_mosaic=args['no_mosaic'], - square_training=args['square_training'] + square_training=args['square_training'], + discard_negative=args["discard_negative"] ) valid_dataset = create_valid_dataset( VALID_DIR_IMAGES, VALID_DIR_LABELS, IMAGE_SIZE, CLASSES, - square_training=args['square_training'] + square_training=args['square_training'], + discard_negative=args["discard_negative"] ) print('Creating data loaders') if args['distributed']: