def csv_from_pascal( file_path_csv: str, images_directory: str, pascal_directory: str, ): bbox_file_ids = [] with open(file_path_csv, "w") as csv_file: csv_file.write( "ImageID,Source,LabelName,Confidence,XMin,XMax,YMin,YMax," "IsOccluded,IsTruncated,IsGroupOf,IsDepiction,IsInside,id,ClassName\n", ) matching_file_ids = matching_ids(pascal_directory, images_directory, ".xml", ".jpg") for file_id in matching_file_ids: bboxes = bounding_boxes_pascal( os.path.join(pascal_directory, file_id + ".xml"), ) if len(bboxes): bbox_file_ids.append(file_id) for bbox in bboxes: csv_file.write( f"{file_id},,,,{bbox['xmin']},{bbox['xmax']}," f"{bbox['ymin']},{bbox['ymax']},,,,,,,{bbox['label']}\n", ) # return list of image file IDs that are included in the CSV return set(bbox_file_ids)
def purge_non_matching( images_dir: str, annotations_dir: str, annotation_format: str, problems_dir: str = None, ) -> Set[str]: """ TODO :param images_dir: :param annotations_dir: :param annotation_format: :param problems_dir: :return: """ # determine the file extensions if annotation_format not in ["darknet", "kitti", "pascal"]: raise ValueError(f"Unsupported annotation format: {annotation_format}") else: annotation_ext = FORMAT_EXTENSIONS[annotation_format] image_ext = ".jpg" # make the problem files directory if necessary, in case it doesn't already exist if problems_dir is not None: os.makedirs(problems_dir, exist_ok=True) # remove files that aren't matches matching_file_ids = matching_ids(annotations_dir, images_dir, annotation_ext, image_ext) for directory in [annotations_dir, images_dir]: for file_name in os.listdir(directory): # only filter out image and Darknet annotation files (this is # needed in case a subdirectory exists in the directory) # and skip the file named "labels.txt" if file_name != "labels.txt" and \ (file_name.endswith(annotation_ext) or file_name.endswith(image_ext)): if os.path.splitext(file_name)[0] not in matching_file_ids: unmatched_file = os.path.join(directory, file_name) if problems_dir is not None: shutil.move(unmatched_file, os.path.join(problems_dir, file_name)) else: os.remove(unmatched_file) return matching_file_ids
def clean_darknet( darknet_dir: str, images_dir: str, label_replacements: Dict, label_removals: List[str] = None, label_keep: List[str] = None, problems_dir: str = None, ): """ TODO :param darknet_dir: :param images_dir: :param label_replacements: :param label_removals: :param problems_dir: :return: """ _logger.info("Cleaning dataset with Darknet annotations") # convert all PNG images to JPG, and remove the original PNG file for file_id in matching_ids(darknet_dir, images_dir, ".txt", ".png"): png_file_path = os.path.join(images_dir, file_id + ".png") png_to_jpg(png_file_path, remove_png=True) # get the set of file IDs of the Darknet-format annotations and corresponding images file_ids = purge_non_matching(images_dir, darknet_dir, "darknet", problems_dir) # loop over all the matching files and clean the Darknet annotations for file_id in tqdm(file_ids): # update the Darknet annotation file src_annotation_file_path = os.path.join(darknet_dir, file_id + ".txt") for line in fileinput.input(src_annotation_file_path, inplace=True): # get the bounding box label parts = line.split() label = parts[0] # skip rewriting this line if it's a label we want removed if (label_removals is not None) and (label in label_removals): continue # skip rewriting this line if it's a label we do not want to keep if (label_keep is not None) and (label not in label_keep): continue # get the bounding box coordinates center_x = float(parts[1]) center_y = float(parts[2]) bbox_width = float(parts[3]) bbox_height = float(parts[4]) if (label_replacements is not None) and (label in label_replacements): # update the label label = label_replacements[label] # make sure we don't have wonky bounding box values # and if so we'll skip them if (center_x > 1.0) or (center_x < 0.0): # report the issue via log message _logger.warning( "Bounding box center X is out of valid range -- skipping " f"in Darknet annotation file {src_annotation_file_path}", ) continue if (center_y > 1.0) or (center_y < 0.0): # report the issue via log message _logger.warning( "Bounding box center Y is out of valid range -- skipping " f"in Darknet annotation file {src_annotation_file_path}", ) continue if (bbox_width > 1.0) or (bbox_width < 0.0): # report the issue via log message _logger.warning( "Bounding box width is out of valid range -- skipping " f"in Darknet annotation file {src_annotation_file_path}", ) continue if (bbox_height > 1.0) or (bbox_height < 0.0): # report the issue via log message _logger.warning( "Bounding box height is out of valid range -- skipping " f"in Darknet annotation file {src_annotation_file_path}", ) continue # write the line back into the file in-place darknet_parts = [ label, f'{center_x:.4f}', f'{center_y:.4f}', f'{bbox_width:.4f}', f'{bbox_height:.4f}', ] print(" ".join(darknet_parts))
def clean_pascal( pascal_dir: str, images_dir: str, label_replacements: Dict = None, label_removals: List[str] = None, label_keep: List[str] = None, problems_dir: str = None, ): """ TODO :param pascal_dir: :param images_dir: :param label_replacements: :param label_removals: :param problems_dir: :return: """ _logger.info("Cleaning dataset with PASCAL annotations") # convert all PNG images to JPG, and remove the original PNG file for file_id in matching_ids(pascal_dir, images_dir, ".xml", ".png"): png_file_path = os.path.join(images_dir, file_id + ".png") png_to_jpg(png_file_path, remove_png=True) # get the set of file IDs of the Darknet-format annotations and corresponding images file_ids = purge_non_matching(images_dir, pascal_dir, "pascal", problems_dir) # loop over all the matching files and clean the PASCAL annotations for i, file_id in tqdm(enumerate(file_ids)): # get the image width and height jpg_file_name = file_id + ".jpg" image_file_path = os.path.join(images_dir, jpg_file_name) image = Image.open(image_file_path) img_width, img_height = image.size # update the image file name in the PASCAL file src_annotation_file_path = os.path.join(pascal_dir, file_id + ".xml") if os.path.exists(src_annotation_file_path): tree = etree.parse(src_annotation_file_path) root = tree.getroot() size = tree.find("size") width = int(size.find("width").text) height = int(size.find("height").text) if (width != img_width) or (height != img_height): # something's amiss that we can't reasonably fix, remove files if problems_dir is not None: for file_path in [ src_annotation_file_path, image_file_path ]: dest_file_path = os.path.join( problems_dir, os.path.split(file_path)[1]) shutil.move(file_path, dest_file_path) else: os.remove(src_annotation_file_path) os.remove(image_file_path) continue # update the image file name file_name = root.find("filename") if (file_name is not None) and (file_name.text != jpg_file_name): file_name.text = jpg_file_name # loop over all bounding boxes for obj in root.iter("object"): # replace all bounding box labels if specified in the replacement dictionary name = obj.find("name") if (name is None) or ((label_removals is not None) and (name.text in label_removals)): # drop the bounding box parent = obj.getparent() parent.remove(obj) # move on, nothing more to do for this box continue # skip rewriting this line if it's a label we do not want to keep elif (name is None) or ((label_keep is not None) and (name.text not in label_keep)): # drop the bounding box parent = obj.getparent() parent.remove(obj) # move on, nothing more to do for this box continue elif (label_replacements is not None) and (name.text in label_replacements): # update the label name.text = label_replacements[name.text] # for each bounding box make sure we have max # values that are one less than the width/height bbox = obj.find("bndbox") bbox_min_x = int(float(bbox.find("xmin").text)) bbox_min_y = int(float(bbox.find("ymin").text)) bbox_max_x = int(float(bbox.find("xmax").text)) bbox_max_y = int(float(bbox.find("ymax").text)) # make sure we don't have wonky values with mins > maxs if (bbox_min_x >= bbox_max_x) or (bbox_min_y >= bbox_max_y): # drop the bounding box _logger.warning( "Dropping bounding box for object in file " f"{src_annotation_file_path} due to invalid " "min/max values", ) parent = obj.getparent() parent.remove(obj) else: # make sure the max values don't go past the edge if bbox_max_x >= img_width: bbox.find("xmax").text = str(img_width - 1) if bbox_max_y >= img_height: bbox.find("ymax").text = str(img_height - 1) # drop the image path, it's not reliable path = root.find("path") if path is not None: parent = path.getparent() parent.remove(path) # drop the image folder, it's not reliable folder = root.find("folder") if folder is not None: parent = folder.getparent() parent.remove(folder) # write the tree back to file tree.write(src_annotation_file_path)
def clean_kitti( kitti_dir: str, images_dir: str, label_replacements: Dict = None, label_removals: List[str] = None, label_keep: List[str] = None, problems_dir: str = None, ): """ TODO :param kitti_dir: :param images_dir: :param label_replacements: :param label_removals: :param problems_dir: :return: """ _logger.info("Cleaning dataset with KITTI annotations") # convert all PNG images to JPG, and remove the original PNG file for file_id in matching_ids(kitti_dir, images_dir, ".txt", ".png"): png_file_path = os.path.join(images_dir, file_id + ".png") png_to_jpg(png_file_path, remove_png=True) # get the set of file IDs of the Darknet-format annotations and corresponding images file_ids = purge_non_matching(images_dir, kitti_dir, "kitti", problems_dir) # loop over all the matching files and clean the KITTI annotations for file_id in tqdm(file_ids): # get the image width and height jpg_file_name = file_id + ".jpg" image_file_path = os.path.join(images_dir, jpg_file_name) image = Image.open(image_file_path) img_width, img_height = image.size # update the image file name in the KITTI annotation file src_annotation_file_path = os.path.join(kitti_dir, file_id + ".txt") for line in fileinput.input(src_annotation_file_path, inplace=True): parts = line.split() label = parts[0] # skip rewriting this line if it's a label we want removed if (label_removals is not None) and (label in label_removals): continue # skip rewriting this line if it's a label we do not want to keep if (label_keep is not None) and (label not in label_keep): continue truncated = parts[1] occluded = parts[2] alpha = parts[3] bbox_min_x = int(float(parts[4])) bbox_min_y = int(float(parts[5])) bbox_max_x = int(float(parts[6])) bbox_max_y = int(float(parts[7])) dim_x = parts[8] dim_y = parts[9] dim_z = parts[10] loc_x = parts[11] loc_y = parts[12] loc_z = parts[13] rotation_y = parts[14] # not all KITTI-formatted files have a score field if len(parts) == 16: score = parts[15] else: score = " " if (label_replacements is not None) and (label in label_replacements): # update the label label = label_replacements[label] # make sure we don't have wonky bounding box values # with mins > maxs, and if so we'll reverse them if bbox_min_x > bbox_max_x: # report the issue via log message _logger.warning( "Bounding box minimum X is greater than the maximum X " f"in KITTI annotation file {src_annotation_file_path}", ) tmp_holder = bbox_min_x bbox_min_x = bbox_max_x bbox_max_x = tmp_holder if bbox_min_y > bbox_max_y: # report the issue via log message _logger.warning( "Bounding box minimum Y is greater than the maximum Y " f"in KITTI annotation file {src_annotation_file_path}", ) tmp_holder = bbox_min_y bbox_min_y = bbox_max_y bbox_max_y = tmp_holder # perform sanity checks on max values if bbox_max_x >= img_width: # report the issue via log message _logger.warning( "Bounding box maximum X is greater than width in KITTI " f"annotation file {src_annotation_file_path}", ) # fix the issue bbox_max_x = img_width - 1 if bbox_max_y >= img_height: # report the issue via log message _logger.warning( "Bounding box maximum Y is greater than height in KITTI " f"annotation file {src_annotation_file_path}", ) # fix the issue bbox_max_y = img_height - 1 # write the line back into the file in-place kitti_parts = [ label, truncated, occluded, alpha, f'{bbox_min_x:.1f}', f'{bbox_min_y:.1f}', f'{bbox_max_x:.1f}', f'{bbox_max_y:.1f}', dim_x, dim_y, dim_z, loc_x, loc_y, loc_z, rotation_y, ] if len(parts) == 16: kitti_parts.append(score) print(" ".join(kitti_parts))
def masked_dataset_to_tfrecords( images_dir: str, masks_dir: str, tfrecord_dir: str, num_shards: int = 1, dataset_base_name: str = "tfrecord", train_pct: float = 1.0, ): """ Creates TFRecord files corresponding to a dataset of JPG images with corresponding set PNG masks. :param images_dir: directory containing image files :param masks_dir: directory containing mask files corresponding to the images :param tfrecord_dir: directory where the output TFRecord files will be written :param num_shards: number of shards :param dataset_base_name: base name of the TFRecord files to be produced :param train_pct: the percentage of images/masks to use for training, with (1.0 minus this value as the validation percentage), if this value is 1.0 then no split will occur """ masks_ext = ".png" images_ext = ".jpg" file_ids = list(matching_ids(masks_dir, images_dir, masks_ext, images_ext)) random.shuffle(file_ids) # create a mapping of base file names and subsets of file IDs if train_pct < 1.0: # get the correct file name prefix for the TFRecord files # based on the presence of a specified file base name tfrecord_file_prefix_train = "train" tfrecord_file_prefix_valid = "valid" if dataset_base_name != "": tfrecord_file_prefix_train = tfrecord_file_prefix_train + "_" + dataset_base_name tfrecord_file_prefix_valid = tfrecord_file_prefix_valid + "_" + dataset_base_name # get the split index to use for splitting into train/valid sets split_index = int(len(file_ids) * train_pct) # map the file prefixes to the sets of file IDs for the split sections split_names_to_ids = { tfrecord_file_prefix_train: file_ids[:split_index], tfrecord_file_prefix_valid: file_ids[split_index:], } # report the number of samples in each split section _logger.info(f"TFRecord dataset contains {len(file_ids[:split_index])} training samples") _logger.info(f"TFRecord dataset contains {len(file_ids[split_index:])} validation samples") else: # we'll just have one base file name mapped to all file IDs if "" == dataset_base_name: tfrecord_file_prefix = "tfrecord" else: tfrecord_file_prefix = dataset_base_name # map the file prefixes to the set of file IDs split_names_to_ids = { tfrecord_file_prefix: file_ids, } # report the number of samples _logger.info(f"TFRecord dataset contains {len(file_ids)} samples (no train/valid split)") # create an iterable of arguments that will be mapped to concurrent future processes args_iterable = [] for base_name, file_ids in split_names_to_ids.items(): num_images = len(file_ids) num_per_shard = int(math.ceil(num_images / num_shards)) for shard_id in range(num_shards): output_filename = os.path.join( tfrecord_dir, f'{base_name}-{str(shard_id).zfill(5)}-of-{str(num_shards).zfill(5)}.tfrecord', ) tfrecord_writing_args = { "output_path": output_filename, "shard_id": shard_id, "num_per_shard": num_per_shard, "num_images": num_images, "file_ids": file_ids, "images_dir": images_dir, "masks_dir": masks_dir, } args_iterable.append(tfrecord_writing_args) # use a ProcessPoolExecutor to facilitate creating the TFRecords in parallel with concurrent.futures.ProcessPoolExecutor() as executor: # map the TFRecord creation function to the iterable of arguments _logger.info(f"Building TFRecords in directory {tfrecord_dir} ") executor.map(_build_write_tfrecord, args_iterable)
def resize_dataset( input_images_dir: str, input_annotations_dir: str, output_images_dir: str, output_annotations_dir: str, new_width: int, new_height: int, annotation_format: str, ): """ Resizes all images and corresponding annotations located within the specified directories. :param input_images_dir: directory where image files are located :param input_annotations_dir: directory where annotation files are located :param output_images_dir: directory where resized image files should be written :param output_annotations_dir: directory where resized annotation files should be written :param new_width: new width to which the image should be resized :param new_height: new height to which the image should be resized :param annotation_format: "coco", "darknet", "kitti", or "pascal" :return: the number of resized image/annotation files """ # only allow for KITTI and PASCAL annotation formats if annotation_format not in ["kitti", "pascal"]: raise ValueError(f"Unsupported annotation format: {annotation_format}") else: annotation_ext = cvdata.common.FORMAT_EXTENSIONS[annotation_format] # create the destination directories in case they don't already exist os.makedirs(output_annotations_dir, exist_ok=True) os.makedirs(output_images_dir, exist_ok=True) resize_arguments_list = [] for image_ext in (".jpg", ".png"): # find matching annotation and image files (i.e. same base file name) file_ids = matching_ids( input_annotations_dir, input_images_dir, annotation_ext, image_ext, ) # loop over all image files and perform scaling/padding on each for file_id in file_ids: resize_arguments = { "file_id": file_id, "image_ext": image_ext, "annotation_format": annotation_format, "annotation_ext": annotation_ext, "input_images_dir": input_images_dir, "input_annotations_dir": input_annotations_dir, "output_images_dir": output_images_dir, "output_annotations_dir": output_annotations_dir, "new_width": new_width, "new_height": new_height, } resize_arguments_list.append(resize_arguments) # use a ProcessPoolExecutor to download the images in parallel with concurrent.futures.ProcessPoolExecutor() as executor: _logger.info("Resizing files") # use the executor to map the download function to the iterable of arguments list(tqdm(executor.map(_resize_image_label, resize_arguments_list), total=len(resize_arguments_list)))
def filter_class_boxes( src_images_dir: str, src_annotations_dir: str, dest_images_dir: str, dest_annotations_dir: str, class_label_counts: Dict, annotation_format: str, darknet_labels_path: str = None, ): """ TODO :param src_images_dir: :param src_annotations_dir: :param dest_images_dir: :param dest_annotations_dir: :param class_label_counts: :param annotation_format: :param darknet_labels_path: path to the labels file corresponding to a Darknet-annotated dataset """ _logger.info( f"Filtering dataset into annotations directory {dest_annotations_dir} " f"and images directory {dest_images_dir}") # make sure we don't have the same directories for src and dest if src_images_dir == dest_images_dir: raise ValueError( "Source and destination image directories are " "the same, must be different", ) if src_annotations_dir == dest_annotations_dir: raise ValueError( "Source and destination annotation directories are " "the same, must be different", ) # determine the file extension to be used for annotations if annotation_format not in ["darknet", "kitti"]: raise ValueError(f"Unsupported annotation format: {annotation_format}") else: annotation_ext = FORMAT_EXTENSIONS[annotation_format] image_ext = ".jpg" # make the destination directories, in case they don't already exist os.makedirs(dest_images_dir, exist_ok=True) os.makedirs(dest_annotations_dir, exist_ok=True) # keep a count of the boxes we've processed for each image class type processed_class_counts = {k: 0 for k in class_label_counts.keys()} # get all file IDs for image/annotation file matches file_ids = \ matching_ids( src_annotations_dir, src_images_dir, annotation_ext, image_ext, ) # only include the labels specified in the class counts dictionary valid_labels = set(class_label_counts.keys()) # if we're processing Darknet annotations then read the labels file to get # a mapping of indices to labels used with the Darknet annotation files darknet_valid_indices = None darknet_index_labels = None if annotation_format == "darknet": # read the Darknet labels into a dictionary mapping label to label index darknet_index_labels = darknet_indices_to_labels(darknet_labels_path) # get the set of valid indices, i.e. all Darknet indices # corresponding to the labels to be included in the filtered dataset darknet_valid_indices = set() for index, darknet_label in darknet_index_labels.items(): if darknet_label in valid_labels: darknet_valid_indices.add(index) # loop over all the possible image/annotation file pairs for file_id in tqdm(file_ids): # only include the file if we find a box for one of the specified labels include_file = False # if any labels are found in the annotation file that aren't included # in the list of image classes to filter then we'll want to remove the # boxes from the annotation file before writing to the destination directory irrelevant_labels_found = False # get the count(s) of boxes per class label annotation_file_name = file_id + annotation_ext src_annotation_path = os.path.join(src_annotations_dir, annotation_file_name) box_counts = _count_boxes(src_annotation_path, annotation_format, darknet_index_labels) for class_label in box_counts.keys(): if class_label in class_label_counts: processed_class_count = processed_class_counts[class_label] if processed_class_counts[class_label] < class_label_counts[ class_label]: include_file = True processed_class_counts[ class_label] = processed_class_count + box_counts[ class_label] else: irrelevant_labels_found = True if include_file: dest_annotation_path = os.path.join(dest_annotations_dir, annotation_file_name) if irrelevant_labels_found: # remove the unnecessary labels or indices from the # annotation and write it into the destination directory _write_with_removed_labels( src_annotation_path, dest_annotation_path, annotation_format, valid_labels, darknet_valid_indices, ) else: # copy te annotation file into the destination directory as-is shutil.copy(src_annotation_path, dest_annotation_path) # copy the source image file into the destination images directory image_file_name = file_id + image_ext src_image_path = os.path.join(src_images_dir, image_file_name) dest_image_path = os.path.join(dest_images_dir, image_file_name) shutil.copy(src_image_path, dest_image_path)
def pascal_to_kitti( pascal_dir: str, images_dir: str, kitti_data_dir: str, kitti_ids_file_name: str = None, move_image_files: bool = False, ) -> int: """ Builds KITTI annotation files from PASCAL VOC annotation XML files. :param pascal_dir: directory containing input PASCAL VOC annotation XML files, all XML files in this directory matching to corresponding JPG files in the images directory will be converted to KITTI format :param images_dir: directory containing image files corresponding to the PASCAL VOC annotation files to be converted to KITTI format, all files matching to PASCAL VOC annotation files will be either copied (default) or moved (if move_image_files is True) to <kitti_data_dir>/image_2 :param kitti_data_dir: directory under which images will be copied or moved into a subdirectory named "image_2" and KITTI annotation files will be written into a subdirectory named "label_2" :param kitti_ids_file_name: file name to contain all file IDs, to be written into the parent directory above the <kitti_data_dir> :param move_image_files: whether or not to move image files to <kitti_data_dir>/image_2 (default is False and image files are copied instead) :return: """ _logger.info(f"Converting from PASCAL to KITTI for images in directory {images_dir}") # create the KITTI directories in case they don't already exist kitti_images_dir = os.path.join(kitti_data_dir, "image_2") kitti_labels_dir = os.path.join(kitti_data_dir, "label_2") for data_dir_name in (kitti_images_dir, kitti_labels_dir): os.makedirs(os.path.join(kitti_data_dir, data_dir_name), exist_ok=True) # assumed file extensions img_ext = ".jpg" pascal_ext = ".xml" # get list of file IDs of the PASCAL VOC annotations and corresponding images file_ids = matching_ids(pascal_dir, images_dir, pascal_ext, img_ext) # write the KITTI IDs file in the KITTI directory's parent directory if kitti_ids_file_name is not None: kitti_ids_file_path = os.path.join(Path(kitti_data_dir).parent, kitti_ids_file_name) with open(kitti_ids_file_path, "w") as kitti_ids_file: for file_id in file_ids: kitti_ids_file.write(f"{file_id}\n") # build KITTI annotations from PASCAL and copy or # move the image files into KITTI images directory conversion_arguments_list = [] for file_id in file_ids: conversion_arguments = { "file_id": file_id, "pascal_ext": pascal_ext, "img_ext": img_ext, "pascal_dir": pascal_dir, "images_dir": images_dir, "kitti_labels_dir": kitti_labels_dir, "kitti_images_dir": kitti_images_dir, "move_image_files": move_image_files, } conversion_arguments_list.append(conversion_arguments) # use a ProcessPoolExecutor to download the images in parallel with concurrent.futures.ProcessPoolExecutor() as executor: # use the executor to map the download function to the iterable of arguments _logger.info(f"Building KITTI labels in directory {kitti_labels_dir} ") list(tqdm(executor.map(single_pascal_to_kitti, conversion_arguments_list), total=len(conversion_arguments_list))) # return the number of annotations converted return len(file_ids)
def kitti_to_darknet( images_dir: str, kitti_dir: str, darknet_dir: str, darknet_labels: str, ): """ Creates equivalent Darknet (YOLO) annotation files corresponding to a dataset with KITTI annotations. :param images_dir: directory containing the dataset's images :param kitti_dir: directory containing the dataset's KITTI annotation files :param darknet_dir: directory where the equivalent Darknet annotation files will be written :param darknet_labels: labels file corresponding to the label indices used in the Darknet annotation files, will be written into the specified Darknet annotations directory """ _logger.info("Converting annotations in KITTI format to Darknet format equivalents") # create the Darknet annotations directory in case it doesn't yet exist os.makedirs(darknet_dir, exist_ok=True) # get list of file IDs of the KITTI annotations and corresponding images annotation_ext = ".txt" image_ext = ".jpg" file_ids = matching_ids(kitti_dir, images_dir, annotation_ext, image_ext) # dictionary of labels to indices label_indices = {} # build Darknet annotations from KITTI for file_id in tqdm(file_ids): # get the image's dimensions image_file_name = file_id + image_ext width, height, _ = image_dimensions(os.path.join(images_dir, image_file_name)) # loop over all annotation lines in the KITTI file and compute Darknet equivalents annotation_file_name = file_id + annotation_ext with open(os.path.join(kitti_dir, annotation_file_name), "r") as kitti_file: darknet_bboxes = [] for line in kitti_file: parts = line.split() label = parts[0] if label in label_indices: label_index = label_indices[label] else: label_index = len(label_indices) label_indices[label] = label_index box_width_pixels = float(parts[6]) - float(parts[4]) + 1 box_height_pixels = float(parts[7]) - float(parts[5]) + 1 darknet_bbox = { "label_index": label_index, "center_x": ((box_width_pixels / 2) + float(parts[4])) / width, "center_y": ((box_height_pixels / 2) + float(parts[5])) / height, "box_width": box_width_pixels / width, "box_height": box_height_pixels / height, } darknet_bboxes.append(darknet_bbox) # write the Darknet annotation boxes into a Darknet annotation file with open(os.path.join(darknet_dir, annotation_file_name), "w") as darknet_file: for darknet_bbox in darknet_bboxes: darknet_file.write( f"{darknet_bbox['label_index']} {darknet_bbox['center_x']} " f"{darknet_bbox['center_y']} {darknet_bbox['box_width']} " f"{darknet_bbox['box_height']}\n" ) # write the Darknet labels into a text file, one label per line, # in order according to the indices used in the annotation files with open(os.path.join(darknet_dir, darknet_labels), "w") as darknet_labels_file: index_labels = {v: k for k, v in label_indices.items()} for i in range(len(index_labels)): darknet_labels_file.write(f"{index_labels[i]}\n")
def _dataset_bbox_examples( images_dir: str, annotations_dir: str, annotation_format: str, darknet_labels: str = None, ) -> pd.DataFrame: """ :param images_dir: directory containing the dataset's *.jpg image files :param annotations_dir: directory containing the dataset's annotation files :param annotation_format: currently supported: "darknet", "kitti", and "pascal" :param darknet_labels: path to the class labels file corresponding to Darknet (YOLO) annotation files, only necessary if using "darknet" annotation format :return: pandas DataFrame with rows corresponding to the dataset's bounding boxes """ # we expect all images to use the *.jpg extension image_ext = ".jpg" # list of bounding box annotations we'll eventually write to CSV bboxes = [] if annotation_format == "pascal": # get the file IDs for all matching image/PASCAL pairs (i.e. the dataset) annotation_ext = ".xml" for file_id in matching_ids( annotations_dir, images_dir, annotation_ext, image_ext, ): # add all bounding boxes from the PASCAL file to the list of boxes pascal_path = os.path.join(annotations_dir, file_id + annotation_ext) tree = ElementTree.parse(pascal_path) root = tree.getroot() for member in root.findall('object'): bbox_values = ( root.find('filename').text, int(root.find('size')[0].text), int(root.find('size')[1].text), member[0].text, int(member[4][0].text), int(member[4][1].text), int(member[4][2].text), int(member[4][3].text), ) bboxes.append(bbox_values) elif annotation_format == "kitti": # get the file IDs for all matching image/KITTI pairs (i.e. the dataset) annotation_ext = ".txt" for file_id in matching_ids( annotations_dir, images_dir, annotation_ext, image_ext, ): # get the image dimensions from the image file since this # info is not present in the corresponding KITTI annotation image_file_name = file_id + image_ext image_path = os.path.join(images_dir, image_file_name) width, height, _ = image_dimensions(image_path) # add all bounding boxes from the KITTI file to the list of boxes kitti_path = os.path.join(annotations_dir, file_id + annotation_ext) with open(kitti_path, "r") as kitti_file: for line in kitti_file: darknet_box = line.split() bbox_values = ( image_file_name, width, height, darknet_box[0], darknet_box[4], darknet_box[5], darknet_box[6], darknet_box[7], ) bboxes.append(bbox_values) elif annotation_format == "darknet": # read class labels into index/label dictionary darknet_index_labels = darknet_indices_to_labels(darknet_labels) # get the file IDs for all matching image/Darknet pairs (i.e. the dataset) annotation_ext = ".txt" file_ids = matching_ids( annotations_dir, images_dir, annotation_ext, image_ext, ) # get the bounding boxes from the annotation files _logger.info("Extracting bounding box info from Darknet annotations...") for file_id in tqdm(file_ids): # get the image dimensions from the image file since this # info is not present in the corresponding KITTI annotation image_file_name = file_id + image_ext image_path = os.path.join(images_dir, image_file_name) width, height, _ = image_dimensions(image_path) # add all bounding boxes from the Darknet file to the list of boxes darknet_path = os.path.join(annotations_dir, file_id + annotation_ext) with open(darknet_path, "r") as darknet_file: for line in darknet_file: darknet_box = line.split() label_index = int(darknet_box[0]) # only use annotations corresponding to the specified labels if label_index not in darknet_index_labels: # skip this annotation line continue center_x = float(darknet_box[1]) * width center_y = float(darknet_box[2]) * height box_width = float(darknet_box[3]) * width box_height = float(darknet_box[4]) * height bbox_values = ( image_file_name, width, height, darknet_index_labels[label_index], int(center_x - (box_width / 2)), int(center_y - (box_height / 2)), int(center_x + (box_width / 2)), int(center_y + (box_height / 2)), ) bboxes.append(bbox_values) else: raise ValueError(f"Unsupported annotation format: {annotation_format}") # stuff the bounding boxes into a pandas DataFrame column_names = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax'] return pd.DataFrame(bboxes, columns=column_names)
def main(): # parse the command line arguments args_parser = argparse.ArgumentParser() args_parser.add_argument( "--annotations", required=True, type=str, help="annotations directory path", ) args_parser.add_argument( "--images", required=False, type=str, help="images directory path", ) args_parser.add_argument( "--format", type=str, required=True, choices=cvdata.common.FORMAT_CHOICES, help="annotation format", ) args_parser.add_argument( "--file_ids", type=str, required=False, help="directory in which to write text files with file IDs for each " "image/annotation in the dataset that contains a label, with " "one file per label", ) args = vars(args_parser.parse_args()) if args["format"] == "tfrecord": # count and report the examples in the collection of TFRecord files examples_count = count_tfrecord_examples(args["annotations"]) print(f"Total number of examples: {examples_count}") else: # the two dictionaries we'll build for final reporting label_counts = {} label_file_ids = {} if args["format"] == "openimages": # read the OpenImages CSV into a pandas DataFrame df_annotations = pd.read_csv(args["annotations"]) df_annotations = df_annotations[['ImageID', 'LabelName']] # TODO get another dataframe from the class descriptions and get the # readable label names from there to map to the LabelName column # whittle it down to only the rows that match to image IDs file_ids = [os.path.splitext(file_name)[0] for file_name in os.listdir(args["images"])] df_annotations = df_annotations[df_annotations["ImageID"].isin(file_ids)] # TODO populate the label counts and label file IDs dictionaries else: annotation_ext = cvdata.common.FORMAT_EXTENSIONS[args["format"]] # only annotations matching to the images are considered to be valid file_ids = matching_ids(args["annotations"], args["images"], annotation_ext, ".jpg") for file_id in file_ids: annotation_file_path = \ os.path.join(args["annotations"], file_id + annotation_ext) # get the images per label count for label, count in count_labels(annotation_file_path, args["format"]).items(): if label in label_counts: label_counts[label] += 1 else: label_counts[label] = 1 # for each label found in the annotation file add this file ID # to the set of file IDs corresponding to the label if args["file_ids"]: if label in label_file_ids: # add this file ID to the existing set for the label label_file_ids[label].add(file_id) else: # first file ID seen for this label so create new set label_file_ids[label] = {file_id} # write the images per label counts for label, count in label_counts.items(): print(f"Label: {label}\t\tCount: {count}") # write the label ID files, if requested if args["file_ids"]: for label, file_ids_for_label in label_file_ids.items(): label_file_ids_path = os.path.join(args["file_ids"], label + ".txt") with open(label_file_ids_path, "w") as label_file_ids_file: for file_id in file_ids_for_label: label_file_ids_file.write(f"{file_id}\n")