xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #226 from hardikdava/fix/yolo_loading_bug 

Fix yolo dataset background instance loading
This commit is contained in:
Piotr Skalski 2023-07-23 14:51:11 +02:00 committed by GitHub
parent 24f013757f 864764e665
commit 97ae155b1e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
supervision/dataset/formats/yolo.py
View File

@ -52,17 +52,6 @@ def _polygons_to_masks(
)
def _pair_image_with_annotation(
image_paths: List[Union[str, Path]], annotation_paths: List[Union[str, Path]]
) -> List[Tuple[Union[str, Path], Union[str, Path]]]:
image_dict = {p.stem: p for p in image_paths}
return [
(image_dict[annotation_path.stem], annotation_path)
for annotation_path in annotation_paths
if annotation_path.stem in image_dict
]
def _with_mask(lines: List[str]) -> bool:
return any([len(line.split()) > 5 for line in lines])
@ -137,14 +126,21 @@ def load_yolo_annotations(
annotation_paths = list_files_with_extensions(
directory=annotations_directory_path, extensions=["txt"]
)
path_pairs = _pair_image_with_annotation(
image_paths=image_paths, annotation_paths=annotation_paths
)
classes = _extract_class_names(file_path=data_yaml_path)
images = {}
annotations = {}
for image_path, annotation_path in path_pairs:
for image_path in image_paths:
image_name = Path(image_path).stem
image = cv2.imread(str(image_path))
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.txt")
if not os.path.exists(annotation_path):
images[image_path.name] = image
annotations[image_path.name] = Detections.empty()
continue
lines = read_txt_file(str(annotation_path))
h, w, _ = image.shape
resolution_wh = (w, h)