diff --git a/embedding_converter/README.md b/embedding_converter/README.md index c088c39..9c1dbc6 100644 --- a/embedding_converter/README.md +++ b/embedding_converter/README.md @@ -26,9 +26,8 @@ Setup This `config.ini` utilizes the MegaFace dataset to train the Embedding Converter for SimSwap. ``` -[preparing.dataset] -dataset_path = .datasets/images -image_pattern = {}/*.*g +[training.dataset] +dataset_file_pattern = .datasets/images/{}/*.*g ``` ``` diff --git a/embedding_converter/config.ini b/embedding_converter/config.ini index 181d835..62504ef 100644 --- a/embedding_converter/config.ini +++ b/embedding_converter/config.ini @@ -1,6 +1,5 @@ -[preparing.dataset] -dataset_path = -image_pattern = +[training.dataset] +dataset_file_pattern = [training.loader] split_ratio = diff --git a/embedding_converter/src/data_loader.py b/embedding_converter/src/data_loader.py index 92723b5..21f335e 100644 --- a/embedding_converter/src/data_loader.py +++ b/embedding_converter/src/data_loader.py @@ -6,12 +6,12 @@ import torchvision.transforms as transforms from torch.utils.data import Dataset from .helper import read_image -from .types import Batch, ImagePathList +from .types import Batch, Paths class DataLoaderRecognition(Dataset[torch.Tensor]): - def __init__(self, dataset_path : str, dataset_image_pattern : str) -> None: - self.image_paths = self.prepare_image_paths(dataset_path, dataset_image_pattern) + def __init__(self, dataset_file_pattern : str) -> None: + self.image_paths = self.prepare_image_paths(dataset_file_pattern) self.dataset_total = len(self.image_paths) self.transforms = self.compose_transforms() @@ -24,9 +24,8 @@ class DataLoaderRecognition(Dataset[torch.Tensor]): def __len__(self) -> int: return self.dataset_total - def prepare_image_paths(self, dataset_path : str, dataset_image_pattern : str) -> ImagePathList: - image_paths = glob.glob(dataset_image_pattern.format(dataset_path)) - return image_paths + def prepare_image_paths(self, dataset_file_pattern : str) -> Paths: + return glob.glob(dataset_file_pattern) def compose_transforms(self) -> transforms: transform = transforms.Compose( diff --git a/embedding_converter/src/helper.py b/embedding_converter/src/helper.py index 50cf44a..c6775f6 100644 --- a/embedding_converter/src/helper.py +++ b/embedding_converter/src/helper.py @@ -1,17 +1,7 @@ -import platform - import cv2 -import numpy from .types import VisionFrame -def is_windows() -> bool: - return platform.system().lower() == 'windows' - - def read_image(image_path : str) -> VisionFrame: - if is_windows(): - image_buffer = numpy.fromfile(image_path, dtype = numpy.uint8) - return cv2.imdecode(image_buffer, cv2.IMREAD_COLOR) return cv2.imread(image_path) diff --git a/embedding_converter/src/training.py b/embedding_converter/src/training.py index 5e8697a..2decc64 100644 --- a/embedding_converter/src/training.py +++ b/embedding_converter/src/training.py @@ -115,11 +115,10 @@ def create_trainer() -> Trainer: def train() -> None: - dataset_path = CONFIG.get('preparing.dataset', 'dataset_path') - dataset_image_pattern = CONFIG.get('preparing.dataset', 'image_pattern') + dataset_file_pattern = CONFIG.get('training.dataset', 'image_pattern') resume_file_path = CONFIG.get('training.output', 'resume_file_path') - dataset = DataLoaderRecognition(dataset_path, dataset_image_pattern) + dataset = DataLoaderRecognition(dataset_file_pattern) training_loader, validation_loader = create_loaders(dataset) embedding_converter_trainer = EmbeddingConverterTrainer() trainer = create_trainer() diff --git a/embedding_converter/src/types.py b/embedding_converter/src/types.py index 78c6ec9..a40d4bb 100644 --- a/embedding_converter/src/types.py +++ b/embedding_converter/src/types.py @@ -3,7 +3,7 @@ from typing import Any, List, TypeAlias from numpy.typing import NDArray from torch import Tensor -ImagePathList : TypeAlias = List[str] +Paths : TypeAlias = List[str] Batch : TypeAlias = Tensor Embedding : TypeAlias = Tensor VisionFrame : TypeAlias = NDArray[Any] diff --git a/requirements.txt b/requirements.txt index 601ba35..36c6ea4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ lightning==2.5.0 numpy==1.26.4 onnx==1.17.0 +onnxruntime==1.20.1 opencv-python==4.11.0.86 pytorch-msssim==1.0.0 torch==2.6.0