Remove editor module

src/training/histogram_dataset.py

@@ -0,0 +1,102 @@
+from torch.utils.data import Dataset
+from typing import List, Optional, Tuple
+from utils import compute_histogram
+from .random_edit import random_edit
+from PIL import Image
+from tqdm import tqdm
+import logging
+import torch
+from pathlib import Path
+
+import PIL.Image
+
+PIL.Image.MAX_IMAGE_PIXELS = None
+
+
+class HistogramDataset(Dataset):
+    def __init__(
+        self,
+        /,
+        paths: List[Path],
+        bin_count: int,
+        edit_count: int = 5,
+        target_size=(240, 240),
+        delete_corrupt_images: bool = False,
+        cache_path: Optional[Path] = None,
+    ):
+        self._paths = sorted(paths)
+        logging.info(f"Loaded {len(self._paths)} original images")
+
+        self._edit_count = edit_count
+        self._bin_count = bin_count
+        self._target_size = target_size
+        self._cache_path = cache_path
+        if self._cache_path:
+            self._cache_path = (
+                self._cache_path
+                / f"{self._bin_count}bins_{self._target_size[0]}x{self._target_size[1]}px"
+            )
+
+        if delete_corrupt_images:
+            self._delete_corrupt_images()
+
+    def _delete_corrupt_images(self) -> None:
+        deleted_count = 0
+        for path in tqdm(self._paths):
+            try:
+                Image.open(path)
+            except:
+                logging.warning(f"Failed to open {path}, deleting...")
+                deleted_count += 1
+                path.unlink()
+        logging.info(f"Deleted {deleted_count} corrupt images")
+
+    def __len__(self):
+        return len(self._paths) * self._edit_count
+
+    def get_original_image(self, original_idx: int) -> Image.Image:
+        original_path = self._paths[original_idx]
+        original = Image.open(original_path)
+        original.thumbnail(
+            self._target_size, Image.Resampling.LANCZOS
+        )  # size will be at most target_size, the aspect ratio is preserved
+        return original
+
+    def get_edited_image(self, original_idx: int, edit_idx: int) -> Image.Image:
+        original_image = self.get_original_image(original_idx)
+        return random_edit(original_image, seed=original_idx * 7919 + edit_idx)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        original_idx = idx // self._edit_count
+        edit_idx = idx % self._edit_count
+
+        if self._cache_path is not None:
+            _cached_data_path = self._cache_path / str(original_idx) / f"{edit_idx}.pt"
+            _cached_data_path.parent.mkdir(parents=True, exist_ok=True)
+            if _cached_data_path.exists():
+                try:
+                    return torch.load(_cached_data_path)
+                except:
+                    logging.warning(
+                        f"Failed to load {_cached_data_path}, regenerating..."
+                    )
+
+        edited = self.get_edited_image(original_idx, edit_idx)
+        edited_histogram = compute_histogram(
+            edited, bins=self._bin_count, normalize=True
+        )
+
+        original = self.get_original_image(original_idx)
+        original_histogram = compute_histogram(
+            original, bins=self._bin_count, normalize=True
+        )
+
+        result = (
+            torch.tensor(edited_histogram, dtype=torch.float).unsqueeze(0),
+            torch.tensor(original_histogram, dtype=torch.float).unsqueeze(0),
+        )
+
+        if self._cache_path is not None:
+            torch.save(result, _cached_data_path)
+
+        return result