refactor: refactor transforms

Author: boomb0om

DPF/transforms/base_file_transforms.py

@@ -5,6 +5,9 @@
 from tqdm.contrib.concurrent import process_map, thread_map
 
 
+PoolOptions = Literal['processes', 'threads']
+
+
 @dataclass
 class TransformsFileData:
  filepath: str
@@ -15,7 +18,7 @@ class BaseFilesTransforms(ABC):
 
  def __init__(
  self,
- pool_type: Literal['processes', 'threads'],
+ pool_type: PoolOptions,
  workers: int = 16,
  pbar: bool = True
  ):
@@ -62,7 +65,7 @@ def data_iterator() -> Iterable[TransformsFileData]:
  arg = TransformsFileData(fp, {k: v[i] for k, v in metadata_lists.items()})
  yield arg
 
- transformed_metadata = pool_map(
+ transformed_metadata: List[TransformsFileData] = pool_map(
  self._process_filepath,
  data_iterator(),
  total=len(paths),

DPF/transforms/image_resize_transforms.py

@@ -2,8 +2,8 @@
 
 from PIL import Image
 
-from DPF.transforms.base_file_transforms import BaseFilesTransforms, TransformsFileData
-from DPF.transforms.image_video_resizer import Resizer
+from DPF.transforms.base_file_transforms import BaseFilesTransforms, TransformsFileData, PoolOptions
+from DPF.transforms.resizer import Resizer
 
 
 class ImageResizeTransforms(BaseFilesTransforms):
@@ -12,7 +12,7 @@ def __init__(
  self,
  resizer: Resizer,
  img_format: str = 'JPEG',
- pool_type: str = 'processes',
+ pool_type: PoolOptions = 'processes',
  workers: int = 16,
  pbar: bool = True
  ):

DPF/transforms/image_video_resizer.py renamed to DPF/transforms/resizer.py

@@ -16,23 +16,19 @@ class Resizer:
 
  def __init__(
  self,
- mode: int,
- fixed_size: Optional[Tuple[int, int]] = None,
- size: Optional[int] = None,
+ mode: ResizerModes,
+ fixed_size: Tuple[int, int] = (512, 512),
+ size: int = 512,
  downscale_only: bool = True
  ):
  self.mode = mode
+ assert self.mode in [ResizerModes.FIXED, ResizerModes.MIN_SIZE, ResizerModes.MAX_SIZE], \
+ f"Invalid resizer mode: {self.mode}. Use mode from ResizerModes"
+
  self.fixed_size = fixed_size
  self.size = size
  self.downscale_only = downscale_only
 
- if self.mode == ResizerModes.FIXED:
- assert len(self.fixed_size) == 2 and all(isinstance(v, int) for v in self.fixed_size)
- elif self.mode in [ResizerModes.MIN_SIZE, ResizerModes.MAX_SIZE]:
- assert isinstance(self.size, int)
- else:
- raise ValueError(f"Invalid resizer mode: {self.mode}. Use mode from ResizerModes")
-
  def get_new_size(self, width: int, height: int) -> Tuple[int, int]:
  if self.mode == ResizerModes.FIXED:
  new_w, new_h = self.fixed_size

DPF/transforms/video_resize_transforms.py

@@ -3,11 +3,11 @@
 import uuid
 from typing import List
 
-from DPF.transforms.base_file_transforms import BaseFilesTransforms, TransformsFileData
-from DPF.transforms.image_video_resizer import Resizer
+from DPF.transforms.base_file_transforms import BaseFilesTransforms, TransformsFileData, PoolOptions
+from DPF.transforms.resizer import Resizer
 
 
-def is_ffmpeg_installed():
+def is_ffmpeg_installed() -> bool:
  try:
  subprocess.run('ffmpeg -version', shell=True, capture_output=True, check=True)
  return True
@@ -21,7 +21,7 @@ def __init__(
  self,
  resizer: Resizer,
  ffmpeg_preset: str = 'fast',
- pool_type: str = 'processes',
+ pool_type: PoolOptions = 'processes',
  workers: int = 16,
  pbar: bool = True
  ):