|
| 1 | +import logging |
| 2 | +import os |
| 3 | +from pathlib import Path |
| 4 | +from typing import Union, Sequence, List, Optional, Dict, Tuple |
| 5 | + |
| 6 | +import PIL.Image |
| 7 | + |
| 8 | +from dagshub_annotation_converter.converters.common import group_annotations_by_filename |
| 9 | +from dagshub_annotation_converter.formats.yolo import ( |
| 10 | + export_lookup, |
| 11 | + allowed_annotation_types, |
| 12 | + YoloContext, |
| 13 | + import_lookup, |
| 14 | + YoloAnnotationTypes, |
| 15 | +) |
| 16 | +from dagshub_annotation_converter.ir.image import IRImageAnnotationBase |
| 17 | +from dagshub_annotation_converter.util import is_image, replace_folder |
| 18 | + |
| 19 | +logger = logging.getLogger(__name__) |
| 20 | + |
| 21 | + |
| 22 | +def load_yolo_from_fs_with_context( |
| 23 | + context: YoloContext, |
| 24 | + import_dir: Union[str, Path] = ".", |
| 25 | +) -> Dict[str, Sequence[IRImageAnnotationBase]]: |
| 26 | + assert context.path is not None |
| 27 | + |
| 28 | + annotations: Dict[str, Sequence[IRImageAnnotationBase]] = {} |
| 29 | + |
| 30 | + import_dir_path = Path(import_dir) |
| 31 | + |
| 32 | + if context.path.is_absolute(): |
| 33 | + data_dir_path = context.path |
| 34 | + else: |
| 35 | + data_dir_path = import_dir_path / context.path |
| 36 | + |
| 37 | + for dirpath, subdirs, files in os.walk(data_dir_path): |
| 38 | + if context.image_dir_name not in dirpath.split("/"): |
| 39 | + logger.debug(f"{dirpath} is not an image dir, skipping") |
| 40 | + continue |
| 41 | + for filename in files: |
| 42 | + fullpath = os.path.join(dirpath, filename) |
| 43 | + img = Path(fullpath) |
| 44 | + relpath = img.relative_to(data_dir_path) |
| 45 | + if not is_image(img): |
| 46 | + logger.debug(f"Skipping {img} because it's not an image") |
| 47 | + continue |
| 48 | + annotation = replace_folder(img, context.image_dir_name, context.label_dir_name, context.label_extension) |
| 49 | + if annotation is None: |
| 50 | + logger.warning(f"Couldn't generate annotation file path for image file [{img}]") |
| 51 | + continue |
| 52 | + if not annotation.exists(): |
| 53 | + logger.warning(f"Couldn't find annotation file [{annotation}] for image file [{img}]") |
| 54 | + continue |
| 55 | + annotations[str(relpath)] = parse_annotation(context, data_dir_path, img, annotation) |
| 56 | + |
| 57 | + return annotations |
| 58 | + |
| 59 | + |
| 60 | +def parse_annotation( |
| 61 | + context: YoloContext, base_path: Path, img_path: Path, annotation_path: Path |
| 62 | +) -> Sequence[IRImageAnnotationBase]: |
| 63 | + img = PIL.Image.open(img_path) |
| 64 | + img_width, img_height = img.size |
| 65 | + |
| 66 | + annotation_strings = annotation_path.read_text().strip().split("\n") |
| 67 | + |
| 68 | + assert context.annotation_type is not None |
| 69 | + |
| 70 | + convert_func = import_lookup[context.annotation_type] |
| 71 | + |
| 72 | + res: List[IRImageAnnotationBase] = [] |
| 73 | + rel_path = str(img_path.relative_to(base_path)) |
| 74 | + |
| 75 | + for ann in annotation_strings: |
| 76 | + res.append(convert_func(ann, context, img_width, img_height, img).with_filename(rel_path)) |
| 77 | + |
| 78 | + return res |
| 79 | + |
| 80 | + |
| 81 | +def load_yolo_from_fs( |
| 82 | + annotation_type: YoloAnnotationTypes, |
| 83 | + meta_file: Union[str, Path] = "annotations.yaml", |
| 84 | + image_dir_name: str = "images", |
| 85 | + label_dir_name: str = "labels", |
| 86 | +) -> Tuple[Dict[str, Sequence[IRImageAnnotationBase]], YoloContext]: |
| 87 | + meta_file_path = Path(meta_file).absolute() |
| 88 | + context = YoloContext.from_yaml_file(meta_file, annotation_type=annotation_type) |
| 89 | + context.image_dir_name = image_dir_name |
| 90 | + context.label_dir_name = label_dir_name |
| 91 | + context.annotation_type = annotation_type |
| 92 | + |
| 93 | + return load_yolo_from_fs_with_context(context, import_dir=meta_file_path.parent), context |
| 94 | + |
| 95 | + |
| 96 | +# ======== Annotation Export ======== # |
| 97 | + |
| 98 | + |
| 99 | +def annotations_to_string(annotations: Sequence[IRImageAnnotationBase], context: YoloContext) -> Optional[str]: |
| 100 | + """ |
| 101 | + Serializes multiple YOLO annotations into the contents of the annotations file. |
| 102 | + Also makes sure that only annotations of the correct type for context.annotation_type are serialized. |
| 103 | +
|
| 104 | + :param annotations: Annotations to serialize (should be single file) |
| 105 | + :param context: Exporting context |
| 106 | + :return: String of the content of the file |
| 107 | + """ |
| 108 | + filtered_annotations = [ |
| 109 | + ann for ann in annotations if isinstance(ann, allowed_annotation_types[context.annotation_type]) |
| 110 | + ] |
| 111 | + |
| 112 | + if len(filtered_annotations) != len(annotations): |
| 113 | + logger.warning( |
| 114 | + f"{annotations[0].filename} has {len(annotations) - len(filtered_annotations)} " |
| 115 | + f"annotations of the wrong type that won't be exported" |
| 116 | + ) |
| 117 | + |
| 118 | + if len(filtered_annotations) == 0: |
| 119 | + return None |
| 120 | + |
| 121 | + export_fn = export_lookup[context.annotation_type] |
| 122 | + |
| 123 | + return "\n".join([export_fn(ann, context) for ann in filtered_annotations]) |
| 124 | + |
| 125 | + |
| 126 | +def export_to_fs( |
| 127 | + context: YoloContext, |
| 128 | + annotations: List[IRImageAnnotationBase], |
| 129 | + export_dir: Union[str, Path] = ".", |
| 130 | + meta_file="yolo_dagshub.yaml", |
| 131 | +) -> Path: |
| 132 | + """ |
| 133 | + Exports annotations to YOLO format. |
| 134 | +
|
| 135 | + This function exports them in a way that allows you to train with YOLO right away, |
| 136 | + as long as the images have already been copied to the data folder. |
| 137 | +
|
| 138 | + :param context: Context for exporting. Set the ``path`` attribute to specify the directory with the data, |
| 139 | + otherwise exports a ``data`` folder in the current working directory. |
| 140 | + :param annotations: Annotations to export |
| 141 | + :param export_dir: Directory to export to. If not specified, exports to the current working directory. |
| 142 | + :param meta_file: Name of the YAML file of the YOLO dataset definition. |
| 143 | + This file will be written to the parent directory of the data path. |
| 144 | +
|
| 145 | + :return: Path to the YAML file with the exported data |
| 146 | + """ |
| 147 | + if context.path is None: |
| 148 | + print(f"`YoloContext.path` was not set. Exporting to {os.path.join(os.getcwd(), 'data')}") |
| 149 | + context.path = Path("data") |
| 150 | + |
| 151 | + grouped_annotations = group_annotations_by_filename(annotations) |
| 152 | + |
| 153 | + export_path = Path(export_dir) |
| 154 | + |
| 155 | + for filename, anns in grouped_annotations.items(): |
| 156 | + annotation_filepath = replace_folder( |
| 157 | + Path(filename), context.image_dir_name, context.label_dir_name, context.label_extension |
| 158 | + ) |
| 159 | + if annotation_filepath is None: |
| 160 | + logger.warning(f"Couldn't generate annotation file path for image file [{filename}]") |
| 161 | + continue |
| 162 | + annotation_filename = export_path / context.path / annotation_filepath |
| 163 | + annotation_filename.parent.mkdir(parents=True, exist_ok=True) |
| 164 | + annotation_content = annotations_to_string(anns, context) |
| 165 | + if annotation_content is not None: |
| 166 | + with open(annotation_filename, "w") as f: |
| 167 | + f.write(annotation_content) |
| 168 | + |
| 169 | + # TODO: test/val splitting |
| 170 | + yaml_file_path = export_path / meta_file |
| 171 | + with open(yaml_file_path, "w") as yaml_f: |
| 172 | + yaml_f.write(context.get_yaml_content()) |
| 173 | + |
| 174 | + logger.warning(f"Saved annotations to {context.path}\nand .YAML file at {yaml_file_path}") |
| 175 | + |
| 176 | + return yaml_file_path.absolute() |
0 commit comments