| from src.DataManager.base import BaseDataManager |
| from src.DataManager.utils import imread_rgb, imwrite_rgb |
|
|
| import numpy as np |
| from pathlib import Path |
|
|
|
|
| class ImageDataManager(BaseDataManager): |
| def __init__(self, src_data: Path, output_dir: Path): |
| self.output_dir: Path = output_dir |
| self.output_dir.mkdir(exist_ok=True) |
| self.output_dir = output_dir / "img" |
| self.output_dir.mkdir(exist_ok=True) |
|
|
| self.data_paths = [] |
| if src_data.is_file(): |
| self.data_paths.append(src_data) |
| elif src_data.is_dir(): |
| self.data_paths = ( |
| list(src_data.glob("*.jpg")) |
| + list(src_data.glob("*.jpeg")) |
| + list(src_data.glob("*.png")) |
| ) |
|
|
| assert len(self.data_paths), "Data must be supplied!" |
|
|
| self.data_paths_iter = iter(self.data_paths) |
|
|
| self.last_idx = -1 |
|
|
| def __len__(self): |
| return len(self.data_paths) |
|
|
| def get(self) -> np.ndarray: |
| img_path = next(self.data_paths_iter) |
| self.last_idx += 1 |
| return imread_rgb(img_path) |
|
|
| def save(self, img: np.ndarray): |
| filename = "swap_" + Path(self.data_paths[self.last_idx]).name |
|
|
| imwrite_rgb(self.output_dir / filename, img) |
|
|
