From f4fd5957832c19a6650affb920c0b23cab42a7c7 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Tue, 9 Jul 2024 14:30:07 +0300 Subject: [PATCH] Fix image poaths in save_dataset, copy images if lazy --- supervision/dataset/utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py index 4da19b62..05e3e081 100644 --- a/supervision/dataset/utils.py +++ b/supervision/dataset/utils.py @@ -3,6 +3,7 @@ import os import random from pathlib import Path from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union +import shutil import cv2 import numpy as np @@ -101,10 +102,13 @@ def save_dataset_images( dataset: "DetectionDataset", images_directory_path: str ) -> None: Path(images_directory_path).mkdir(parents=True, exist_ok=True) - for image_path, image, _ in dataset: - final_path = os.path.join(images_directory_path, image_path) - cv2.imwrite(final_path, image) - + for image_path in dataset.image_paths: + final_path = os.path.join(images_directory_path, Path(image_path).name) + if image_path in dataset._images_in_memory: + image = dataset._images_in_memory[image_path] + cv2.imwrite(final_path, image) + else: + shutil.copyfile(image_path, final_path) def train_test_split( data: List[T],