From 2c366e976c2a024553d8eb231345efedee18c661 Mon Sep 17 00:00:00 2001 From: zeptofine Date: Wed, 18 Oct 2023 11:27:48 -0400 Subject: [PATCH] Enums, more resize flexibility --- imdataset_creator/datarules/base_rules.py | 2 +- .../datarules/dataset_builder.py | 9 +- imdataset_creator/datarules/image_rules.py | 1 - imdataset_creator/file.py | 39 ++++-- imdataset_creator/gui/output_filters.py | 45 +++++-- imdataset_creator/gui/rule_views.py | 10 +- imdataset_creator/image_filters/destroyers.py | 124 +++++++++++------- imdataset_creator/image_filters/resizer.py | 57 +++++--- pyproject.toml | 1 - 9 files changed, 190 insertions(+), 98 deletions(-) diff --git a/imdataset_creator/datarules/base_rules.py b/imdataset_creator/datarules/base_rules.py index 6529360..4fb3ffa 100644 --- a/imdataset_creator/datarules/base_rules.py +++ b/imdataset_creator/datarules/base_rules.py @@ -167,7 +167,7 @@ def get_field(self, field_name: str, args: Sequence[Any], kwargs: Mapping[str, A DEFAULT_OUTPUT_FORMAT = "{relative_path}/{file}.{ext}" -PLACEHOLDER_FORMAT_FILE = File("/folder/subfolder/to/file.png", "/folder", "subfolder/to", "file", ".png") +PLACEHOLDER_FORMAT_FILE = File.from_src(Path("/folder"), Path("/folder/subfolder/to/file.png")) PLACEHOLDER_FORMAT_KWARGS = PLACEHOLDER_FORMAT_FILE.to_dict() diff --git a/imdataset_creator/datarules/dataset_builder.py b/imdataset_creator/datarules/dataset_builder.py index 1f737f6..6d0a235 100644 --- a/imdataset_creator/datarules/dataset_builder.py +++ b/imdataset_creator/datarules/dataset_builder.py @@ -8,19 +8,14 @@ from typing import Generator, Literal, TypeVar, overload import polars as pl -from polars import DataFrame, Expr, Series +from polars import DataFrame, Expr from polars.type_aliases import SchemaDefinition -from ..configs import MainConfig -from ..file import File -from ..scenarios import FileScenario, OutputScenario from .base_rules import ( Comparable, DataTypeSchema, ExprDict, FastComparable, - Input, - Output, Producer, ProducerSchema, ProducerSet, @@ -242,6 +237,7 @@ def populate_chunks( db_schema = self.type_schema chunk: DataFrame for chunk in chunks: + # current_paths = list(chunk.get_column("path")) # used for debugging for schema in schemas: chunk = chunk.with_columns(**schema) chunk = chunk.select(db_schema) @@ -332,4 +328,3 @@ def comply_to_schema(self, schema: SchemaDefinition, in_place: bool = False) -> if in_place: self.__df = new_df return new_df - diff --git a/imdataset_creator/datarules/image_rules.py b/imdataset_creator/datarules/image_rules.py index adecbfc..90b06f8 100644 --- a/imdataset_creator/datarules/image_rules.py +++ b/imdataset_creator/datarules/image_rules.py @@ -147,7 +147,6 @@ def __call__(self) -> ProducerSchema: return [{"hash": col("path").apply(self._hash_img)}] def _hash_img(self, pth) -> str: - assert self.hasher is not None return str(self.hasher(Image.open(pth))) diff --git a/imdataset_creator/file.py b/imdataset_creator/file.py index 9981d8b..33cae61 100644 --- a/imdataset_creator/file.py +++ b/imdataset_creator/file.py @@ -1,13 +1,36 @@ +import os from dataclasses import dataclass from pathlib import Path +from .configs.keyworded import fancy_repr + +class MalleablePath(str): + def __format__(self, format_spec): + formats = format_spec.split(",") + newfmt: MalleablePath = self + for fmt in formats: + if "=" in fmt: + key, val = fmt.split("=") + if key == "maxlen": + newfmt = MalleablePath(newfmt[: int(val)]) + else: + raise ValueError(f"Unknown format specifier: {key}") + elif fmt == "underscores": + newfmt = MalleablePath("_".join(self.split(" "))) + elif fmt == "underscore_path": + newfmt = MalleablePath("_".join(Path(self).parts)) + + return str(newfmt) + + +@fancy_repr @dataclass(frozen=True) class File: - absolute_pth: str - src: str - relative_path: str - file: str + absolute_pth: MalleablePath + src: MalleablePath + relative_path: MalleablePath + file: MalleablePath ext: str def to_dict(self): @@ -22,9 +45,9 @@ def to_dict(self): @classmethod def from_src(cls, src: Path, pth: Path): return cls( - absolute_pth=str(pth.resolve()), - src=str(src), - relative_path=str(pth.relative_to(src).parent), - file=pth.stem, + absolute_pth=MalleablePath(pth.resolve()), + src=MalleablePath(src), + relative_path=MalleablePath(pth.relative_to(src).parent), + file=MalleablePath(pth.stem), ext=pth.suffix[1:], ) diff --git a/imdataset_creator/gui/output_filters.py b/imdataset_creator/gui/output_filters.py index 6edb953..89fe108 100644 --- a/imdataset_creator/gui/output_filters.py +++ b/imdataset_creator/gui/output_filters.py @@ -1,5 +1,18 @@ from PySide6.QtCore import QSize -from PySide6.QtWidgets import QDoubleSpinBox, QLabel, QSpinBox +from PySide6.QtWidgets import ( + QCheckBox, + QComboBox, + QDoubleSpinBox, + QFrame, + QLabel, + QLineEdit, + QProgressBar, + QPushButton, + QSlider, + QSpinBox, + QToolButton, + QWidget, +) from ..datarules.base_rules import Filter from ..image_filters import destroyers, resizer @@ -34,13 +47,16 @@ class ResizeFilterView(FilterView): bound_item = resizer.Resize def configure_settings_group(self): + self.resize_mode = QComboBox(self) + self.resize_mode.addItems(resizer.ResizeMode._member_names_) self.scale = QDoubleSpinBox(self) - self.scale.setSuffix("%") self.scale.setMinimum(1) - self.scale.setMaximum(1_000) + self.scale.setMaximum(100_000) - self.group_grid.addWidget(QLabel("Scale:", self), 0, 0) - self.group_grid.addWidget(self.scale, 0, 1) + self.group_grid.addWidget(QLabel("Resize mode: ", self), 0, 0) + self.group_grid.addWidget(self.resize_mode, 0, 1) + self.group_grid.addWidget(QLabel("Scale:", self), 1, 0) + self.group_grid.addWidget(self.scale, 1, 1) def reset_settings_group(self): self.scale.setValue(100) @@ -89,6 +105,15 @@ def get_config(self) -> resizer.CropData: "height": val if (val := self.height_box.value()) else None, } + @classmethod + def from_config(cls, cfg: resizer.CropData, parent=None): + self = cls(parent) + self.left_box.setValue(cfg["left"] or 0) + self.top_box.setValue(cfg["top"] or 0) + self.width_box.setValue(cfg["width"] or 0) + self.height_box.setValue(cfg["height"] or 0) + return self + class BlurFilterView(FilterView): title = "Blur" @@ -97,7 +122,7 @@ class BlurFilterView(FilterView): bound_item = destroyers.Blur def configure_settings_group(self): - self.algorithms = MiniCheckList(destroyers.AllBlurAlgos, self) + self.algorithms = MiniCheckList(destroyers.BlurAlgorithm._member_names_, self) self.scale = QDoubleSpinBox(self) scale_label = QLabel("Scale:", self) tooltip(scale_label, TOOLTIPS["blur_scale"]) @@ -134,14 +159,14 @@ def get_config(self) -> destroyers.BlurData: { "algorithms": algos, "blur_range": [self.blur_range_x.value(), self.blur_range_y.value()], - "scale": self.scale.value() / 100, + "scale": self.scale.value(), } ) @classmethod def from_config(cls, cfg, parent=None): self = cls(parent) - self.scale.setValue(cfg["scale"] * 100) + self.scale.setValue(cfg["scale"]) for item in cfg["algorithms"]: self.algorithms.set_config(item, True) r_x, r_y = cfg["blur_range"] @@ -158,7 +183,7 @@ class NoiseFilterView(FilterView): bound_item = destroyers.Noise def configure_settings_group(self): - self.algorithms = MiniCheckList(destroyers.AllNoiseAlgos, self) + self.algorithms = MiniCheckList(destroyers.NoiseAlgorithm._member_names_, self) self.scale = QDoubleSpinBox(self) self.scale.setSuffix("%") self.scale.setMinimum(1) @@ -215,7 +240,7 @@ class CompressionFilterView(FilterView): bound_item = destroyers.Compression def configure_settings_group(self): - self.algorithms = MiniCheckList(destroyers.AllCompressionAlgos, self) + self.algorithms = MiniCheckList(destroyers.CompressionAlgorithms._member_names_, self) self.group_grid.addWidget(self.algorithms, 0, 0, 1, 3) # jpeg quality diff --git a/imdataset_creator/gui/rule_views.py b/imdataset_creator/gui/rule_views.py index 5551b65..f507ba6 100644 --- a/imdataset_creator/gui/rule_views.py +++ b/imdataset_creator/gui/rule_views.py @@ -54,9 +54,12 @@ def get_config_wrapper(self: RuleView): def set_requires(self, val): newdesc = self.__original_desc if val: - newdesc += f"\n requires: {val}" + newdesc = newdesc + ("\n" if newdesc else "") + f"requires: {val}" print("updated requires") + self.desc = newdesc self.description_widget.setText(newdesc) + if not self.description_widget.isVisible(): + self.description_widget.show() class ItemsUnusedError(ValueError): @@ -157,8 +160,8 @@ def reset_settings_group(self): def get(self): super().get() return data_rules.BlackWhitelistRule( - self.whitelist.toPlainText().splitlines(), - self.blacklist.toPlainText().splitlines(), + whitelist=self.whitelist.toPlainText().splitlines(), + blacklist=self.blacklist.toPlainText().splitlines(), ) def get_config(self) -> data_rules.BlackWhitelistData: @@ -238,7 +241,6 @@ def reset_settings_group(self): self.crop.setChecked(True) def get(self): - super().get() return image_rules.ResRule( min_res=self.min.value(), max_res=self.max.value(), diff --git a/imdataset_creator/image_filters/destroyers.py b/imdataset_creator/image_filters/destroyers.py index 450e3c3..0d12b4a 100644 --- a/imdataset_creator/image_filters/destroyers.py +++ b/imdataset_creator/image_filters/destroyers.py @@ -2,9 +2,10 @@ import subprocess import typing from dataclasses import dataclass +from enum import Enum from math import sqrt from random import choice, randint -from typing import Literal +from typing import Literal, Self import cv2 import ffmpeg @@ -19,8 +20,11 @@ np_gen = np.random.default_rng() -BlurAlgorithms = Literal["average", "gaussian", "isotropic", "anisotropic"] -AllBlurAlgos = typing.get_args(BlurAlgorithms) +class BlurAlgorithm(Enum): + AVERAGE = 0 + GAUSSIAN = 1 + ISOTROPIC = 2 + ANISOTROPIC = 3 class BlurData(FilterData): @@ -31,7 +35,7 @@ class BlurData(FilterData): @dataclass(frozen=True) class Blur(Filter): - algorithms: list[BlurAlgorithms] | None = None + algorithms: list[BlurAlgorithm] | None = None blur_range: tuple[int, int] = (1, 16) scale: float = 0.25 @@ -39,34 +43,47 @@ def run( self, img: np.ndarray, ) -> np.ndarray: - algorithms = self.algorithms or ["average"] - - algorithm: BlurAlgorithms = choice(algorithms) + algorithms = self.algorithms or [BlurAlgorithm.AVERAGE] + algorithm: BlurAlgorithm = choice(algorithms) start, stop = self.blur_range - if algorithm in ["average", "gaussian"]: - ksize: int = int((randint(start, stop) | 1) * self.scale) + ri = randint(start, stop) + ksize: int + if algorithm == BlurAlgorithm.AVERAGE: + ksize = int(ri * self.scale) + ksize = ksize + (ksize % 2 == 0) # ensure ksize is odd + return cv2.blur(img, (ksize, ksize)) + if algorithm == BlurAlgorithm.GAUSSIAN: + ksize = int((ri | 1) * self.scale) ksize = ksize + (ksize % 2 == 0) # ensure ksize is odd - - if algorithm == "average": - return cv2.blur(img, (ksize, ksize)) return cv2.GaussianBlur(img, (ksize, ksize), 0) - if algorithm in ["isotropic", "anisotropic"]: - sigma1: float = randint(start, stop) * self.scale + if algorithm == BlurAlgorithm.ISOTROPIC or algorithm == BlurAlgorithm.ANISOTROPIC: + sigma1: float = ri * self.scale ksize1: int = 2 * int(4 * sigma1 + 0.5) + 1 - if algorithm == "anisotropic": + if algorithm == BlurAlgorithm.ANISOTROPIC: return cv2.GaussianBlur(img, (ksize1, ksize1), sigmaX=sigma1, sigmaY=sigma1) - sigma2: float = randint(start, stop) * self.scale + sigma2: float = ri * self.scale ksize2: int = 2 * int(4 * sigma2 + 0.5) + 1 return cv2.GaussianBlur(img, (ksize1, ksize2), sigmaX=sigma1, sigmaY=sigma2) return img + @classmethod + def from_cfg(cls, cfg: BlurData) -> Self: + return cls( + algorithms=[BlurAlgorithm._member_map_[k] for k in cfg["algorithms"]], # type: ignore + blur_range=cfg["blur_range"], # type: ignore + scale=cfg["scale"], + ) + -NoiseAlgorithms = Literal["uniform", "gaussian", "color", "gray"] -AllNoiseAlgos = typing.get_args(NoiseAlgorithms) +class NoiseAlgorithm(Enum): + UNIFORM = 0 + GAUSSIAN = 1 + COLOR = 2 + GRAY = 3 class NoiseData(FilterData): @@ -77,25 +94,25 @@ class NoiseData(FilterData): @dataclass(frozen=True) class Noise(Filter): - algorithms: list[NoiseAlgorithms] | None = None + algorithms: list[NoiseAlgorithm] | None = None intensity_range: tuple[int, int] = (1, 16) scale: float = 0.25 def run(self, img: ndarray) -> ndarray: - algorithms = self.algorithms or ["uniform"] + algorithms = self.algorithms or [NoiseAlgorithm.UNIFORM] algorithm = choice(algorithms) - if algorithm == "uniform": + if algorithm == NoiseAlgorithm.UNIFORM: intensity = randint(*self.intensity_range) * self.scale noise = np_gen.uniform(-intensity, intensity, img.shape) return cv2.add(img, noise.astype(img.dtype)) - if algorithm == "gaussian": + if algorithm == NoiseAlgorithm.GAUSSIAN: sigma = sqrt(randint(*self.intensity_range) * self.scale) noise = np_gen.normal(0, sigma, img.shape) return cv2.add(img, noise.astype(img.dtype)) - if algorithm == "color": + if algorithm == NoiseAlgorithm.COLOR: noise = np.zeros_like(img) s = (randint(*self.intensity_range), randint(*self.intensity_range), randint(*self.intensity_range)) cv2.randn(noise, 0, s) # type: ignore @@ -108,15 +125,13 @@ def run(self, img: ndarray) -> ndarray: return img + noise -CompressionAlgorithms = Literal[ - "jpeg", - "webp", - "h264", - "hevc", - "mpeg", - "mpeg2", -] -AllCompressionAlgos: tuple[str, ...] = typing.get_args(CompressionAlgorithms) +class CompressionAlgorithms(Enum): + JPEG = "jpeg" + WEBP = "webp" + H264 = "h264" + HEVC = "hevc" + MPEG = "mpeg" + MPEG2 = "mpeg2" class CompressionData(FilterData): @@ -140,35 +155,40 @@ class Compression(Filter): mpeg2_bitrate: int = 4_000_000 def run(self, img: ndarray): - algos = self.algorithms or ["jpeg"] + algos = self.algorithms or [CompressionAlgorithms.JPEG] algorithm = choice(algos) quality: int enc_img: ndarray - if algorithm == "jpeg": + if algorithm == CompressionAlgorithms.JPEG: quality = randint(*self.jpeg_quality_range) enc_img = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), quality])[1] return cv2.imdecode(enc_img, 1) - if algorithm == "webp": + if algorithm == CompressionAlgorithms.WEBP: quality = randint(*self.webp_quality_range) enc_img = cv2.imencode(".webp", img, [int(cv2.IMWRITE_WEBP_QUALITY), quality])[1] return cv2.imdecode(enc_img, 1) - if algorithm in ["h264", "hevc", "mpeg", "mpeg2"]: + if algorithm in [ + CompressionAlgorithms.H264, + CompressionAlgorithms.HEVC, + CompressionAlgorithms.MPEG, + CompressionAlgorithms.MPEG2, + ]: height, width, _ = img.shape - codec = algorithm + codec = algorithm.value container = "mpeg" output_args: dict[str, int | str] crf: int - if algorithm == "h264": + if algorithm == CompressionAlgorithms.H264: crf = randint(*self.h264_crf_range) output_args = {"crf": crf} - elif algorithm == "hevc": + elif algorithm == CompressionAlgorithms.HEVC: crf = randint(*self.hevc_crf_range) output_args = {"crf": crf, "x265-params": "log-level=0"} - elif algorithm == "mpeg": + elif algorithm == CompressionAlgorithms.MPEG: output_args = {"b": self.mpeg_bitrate} codec = "mpeg1video" @@ -195,14 +215,26 @@ def run(self, img: ndarray): try: compressor.wait(10) + newimg = np.frombuffer(out, np.uint8) + if len(newimg) != height * width * 3: + log.warning("New image size does not match") + newimg = newimg[: height * width * 3] # idrk why i need this sometimes + + return newimg.reshape((height, width, 3)) except subprocess.TimeoutExpired as e: compressor.send_signal("SIGINT") log.warning(f"{e}") - newimg = np.frombuffer(out, np.uint8) - if len(newimg) != height * width * 3: - log.warning("New image size does not match") - newimg = newimg[: height * width * 3] # idrk why i need this sometimes - - return newimg.reshape((height, width, 3)) return img + + @classmethod + def from_cfg(cls, cfg: CompressionData) -> Self: + return cls( + algorithms=[CompressionAlgorithms._member_map_[k] for k in cfg["algorithms"]], # type: ignore + jpeg_quality_range=cfg["jpeg_quality_range"], # type: ignore + webp_quality_range=cfg["webp_quality_range"], # type: ignore + h264_crf_range=cfg["h264_crf_range"], # type: ignore + hevc_crf_range=cfg["hevc_crf_range"], # type: ignore + mpeg_bitrate=cfg["mpeg_bitrate"], + mpeg2_bitrate=cfg["mpeg2_bitrate"], + ) diff --git a/imdataset_creator/image_filters/resizer.py b/imdataset_creator/image_filters/resizer.py index 84a6563..82e6188 100644 --- a/imdataset_creator/image_filters/resizer.py +++ b/imdataset_creator/image_filters/resizer.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from enum import Enum from random import choice from typing import Literal @@ -11,20 +12,29 @@ np_gen = np.random.default_rng() -RESIZE_ALGOS = { - "bicubic": cv2.INTER_CUBIC, - "bilinear": cv2.INTER_LINEAR, - "box": cv2.INTER_AREA, - "nearest": cv2.INTER_NEAREST, - "lanczos": cv2.INTER_LANCZOS4, -} -ResizeAlgorithms = Literal["bicubic", "bilinear", "box", "nearest", "lanczos", "down_up"] -DownUpScaleAlgorithms = list(RESIZE_ALGOS.keys()) + +class ResizeAlgos(Enum): + BICUBIC = cv2.INTER_CUBIC + BILINEAR = cv2.INTER_LINEAR + BOX = cv2.INTER_AREA + NEAREST = cv2.INTER_NEAREST + LANCZOS = cv2.INTER_LANCZOS4 + DOWN_UP = False + + +DownUpAlgos = [e for e in ResizeAlgos if e != ResizeAlgos.DOWN_UP] + + +class ResizeMode(Enum): + VALUE = 0 + MAX_RESOLUTION = 1 + MIN_RESOLUTION = 2 @dataclass(frozen=True) class Resize(Filter): - algorithms: list[ResizeAlgorithms] | None = None + mode: ResizeMode = ResizeMode.MIN_RESOLUTION + algorithms: list[ResizeAlgos] | None = None down_up_range: tuple[float, float] = (0.5, 2) scale: float = 0.5 @@ -35,19 +45,26 @@ def run( algorithms = self.algorithms original_algos = algorithms if algorithms is None: - algorithms = ["down_up"] + algorithms = [ResizeAlgos.DOWN_UP] algorithm = choice(algorithms) h, w = img.shape[:2] - new_h = int(h * self.scale) - new_w = int(w * self.scale) + scale = self.scale + if self.mode == ResizeMode.MAX_RESOLUTION: + scale = min(1, self.scale / max(h, w)) + elif self.mode == ResizeMode.MIN_RESOLUTION: + scale = max(1, self.scale / min(h, w)) + + new_h = int(h * scale) + new_w = int(w * scale) + if algorithm == "down_up": - algo1: str - algo2: str + algo1: ResizeAlgos + algo2: ResizeAlgos if original_algos is None: - algo1 = choice(DownUpScaleAlgorithms) - algo2 = choice(DownUpScaleAlgorithms) + algo1 = choice(DownUpAlgos) + algo2 = choice(DownUpAlgos) else: algo1 = original_algos[0] algo2 = original_algos[-1] @@ -57,16 +74,16 @@ def run( cv2.resize( img, (int(w * scale_factor), int(h * scale_factor)), - interpolation=RESIZE_ALGOS[algo1], + interpolation=algo1.value, ), (new_w, new_h), - interpolation=RESIZE_ALGOS[algo2], + interpolation=algo2.value, ) return cv2.resize( img, (new_w, new_h), - interpolation=RESIZE_ALGOS[algorithm], + interpolation=algorithm.value, ) diff --git a/pyproject.toml b/pyproject.toml index b57ecd8..aa918c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,6 @@ extend-select = [ "F", "RET", "SIM", - "TRY", "NPY", "PERF", "RUF",