WIP

data/.gitignore

@@ -0,0 +1 @@
+fivek/

eiuie/base_model.py

@@ -0,0 +1,113 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+
+def BGR2HSI(image: np.ndarray) -> np.ndarray:
+    """
+    Convert image from BGR to HSI.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        Image to be converted, as BGR.
+
+    Returns
+    -------
+    np.ndarray
+        Converted image, as HSI.
+    """
+
+    # Normalize BGR values to [0,1]
+    bgr = image.astype(np.float32) / 255.0
+    blue = bgr[:, :, 0]
+    green = bgr[:, :, 1]
+    red = bgr[:, :, 2]
+
+    # Compute intensity
+    I = (blue + green + red) / 3.0
+
+    # Compute saturation
+    min_val = np.minimum(np.minimum(blue, green), red)
+    S = 1 - 3.0 / (blue + green + red + 1e-6) * min_val
+
+    # Compute hue
+    num = 0.5 * ((red - green) + (red - blue))
+    den = np.sqrt((red - green) ** 2 + (red - blue) * (green - blue))
+    theta = np.arccos(num / (den + 1e-6))
+    H = theta
+    H[blue > green] = 2 * np.pi - H[blue > green]
+    H /= 2 * np.pi
+
+    hsi = np.stack([H, S, I], axis=2)
+    return hsi
+
+
+def HSI2BGR(image: np.ndarray) -> np.ndarray:
+    """
+    Convert image from HSI to BGR.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        Image to be converted, as HSI.
+
+    Returns
+    -------
+    np.ndarray
+        Converted image, as BGR.
+    """
+
+    H = image[:, :, 0] * 2 * np.pi
+    S = image[:, :, 1]
+    I = image[:, :, 2]
+
+    R = np.zeros_like(H)
+    G = np.zeros_like(H)
+    B = np.zeros_like(H)
+
+    # RG sector
+    cond = np.logical_and(0 <= H, H < 2 * np.pi / 3)
+    B[cond] = I[cond] * (1 - S[cond])
+    R[cond] = I[cond] * (1 + S[cond] * np.cos(H[cond]) / np.cos(np.pi / 3 - H[cond]))
+    G[cond] = 3 * I[cond] - (R[cond] + B[cond])
+
+    # GB sector
+    cond = np.logical_and(2 * np.pi / 3 <= H, H < 4 * np.pi / 3)
+    H[cond] = H[cond] - 2 * np.pi / 3
+    R[cond] = I[cond] * (1 - S[cond])
+    G[cond] = I[cond] * (1 + S[cond] * np.cos(H[cond]) / np.cos(np.pi / 3 - H[cond]))
+    B[cond] = 3 * I[cond] - (R[cond] + G[cond])
+
+    # BR sector
+    cond = np.logical_and(4 * np.pi / 3 <= H, H < 2 * np.pi)
+    H[cond] = H[cond] - 4 * np.pi / 3
+    G[cond] = I[cond] * (1 - S[cond])
+    B[cond] = I[cond] * (1 + S[cond] * np.cos(H[cond]) / np.cos(np.pi / 3 - H[cond]))
+    R[cond] = 3 * I[cond] - (G[cond] + B[cond])
+
+    bgr = np.stack([B, G, R], axis=2)
+    return (bgr * 255).astype(np.uint8)
+
+
+class BaseModel(ABC):
+    """
+    Abstract class for all models.
+    """
+
+    @abstractmethod
+    def process_image(self, image: np.ndarray) -> np.ndarray:
+        """
+        Process image using the model.
+
+        Parameters
+        ----------
+        image : np.ndarray
+            Image to be processed, as BGR.
+
+        Returns
+        -------
+        np.ndarray
+            Processed image, as BGR.
+        """
+        ...

eiuie/download.py

@@ -0,0 +1,91 @@
+import os
+import requests
+from pathlib import Path
+import rawpy
+from PIL import Image
+from tqdm import tqdm
+from joblib import Parallel
+
+
+class ProgressParallel(Parallel):
+    def __init__(self, use_tqdm=True, total=None, *args, **kwargs):
+        self._use_tqdm = use_tqdm
+        self._total = total
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        with tqdm(disable=not self._use_tqdm, total=self._total) as self._pbar:
+            return Parallel.__call__(self, *args, **kwargs)
+
+    def print_progress(self):
+        if self._total is None:
+            self._pbar.total = self.n_dispatched_tasks
+        self._pbar.n = self.n_completed_tasks
+        self._pbar.refresh()
+
+
+def dng_to_jpg(img_path: str):
+    raw = rawpy.imread(img_path)
+    rgb = raw.postprocess(use_camera_wb=True)
+    new_name = img_path[:-3] + "jpg"
+    Image.fromarray(rgb).save(new_name, optimize=True)
+    os.remove(img_path)
+
+
+def tif_to_jpg(img_path: str):
+    img = Image.open(img_path)
+    img.save(img_path[:-3] + "jpg")
+    os.remove(img_path)
+
+
+def download_img(img_name: str, data_dir: Path):
+    try:
+        dng_url = f"https://data.csail.mit.edu/graphics/fivek/img/dng/{img_name}.dng"
+        dng_file = requests.get(dng_url, allow_redirects=True)
+        dng_path = data_dir / "raw" / (img_name + ".dng")
+        open(dng_path, "wb").write(dng_file.content)
+        dng_to_jpg(str(dng_path))
+
+        tif_base = "https://data.csail.mit.edu/graphics/fivek/img/tiff16"
+        for expert in ["a", "b", "c", "d", "e"]:
+            url = f"{tif_base}_{expert}/{img_name}.tif"
+            path = data_dir / expert / (img_name + ".tif")
+            tif_file = requests.get(url, allow_redirects=True)
+            open(path, "wb").write(tif_file.content)
+            tif_to_jpg(str(path))
+
+    except Exception as e:
+        print(img_name, e)
+
+
+def download_dataset(store_dir: Path, n_jobs: int = 8):
+    # * Create folders
+    dng_dir = store_dir / "raw"
+    tif_dirs = [store_dir / s for s in ["a", "b", "c", "d", "e"]]
+
+    dng_dir.mkdir(parents=True, exist_ok=True)
+    for path in tif_dirs:
+        path.mkdir(parents=True, exist_ok=True)
+
+    # * Get image info
+    f1 = requests.get(
+        "https://data.csail.mit.edu/graphics/fivek/legal/filesAdobe.txt"
+    ).text.split("\n")
+    f2 = requests.get(
+        "https://data.csail.mit.edu/graphics/fivek/legal/filesAdobeMIT.txt"
+    ).text.split("\n")
+    names = [x for x in set(f1 + f2) if x != ""]
+
+    # * Download imgs
+    ProgressParallel(n_jobs=n_jobs, total=len(names))(
+        download_img(name, store_dir) for name in names
+    )
+
+
+def main():
+    store_dir = Path("data/fivek")
+    download_dataset(store_dir, os.cpu_count() or 1)
+
+
+if __name__ == "__main__":
+    main()

eiuie/homomorphic_filtering.py

@@ -0,0 +1,99 @@
+import numpy as np
+import cv2
+
+import base_model as bm
+
+
+class HomomorphicFiltering(bm.BaseModel):
+    """Homomorphic Filtering"""
+
+    ksize: int
+    sigma: float
+    gamma_1: float
+    gamma_2: float
+    rho: float
+
+    def __init__(self, ksize: int = 3, sigma: float = 1.0):
+        """
+        Parameters
+        ----------
+        ksize : int, optional
+            Kernel size, by default 3
+        sigma : float, optional
+            Gaussian kernel standard deviation, by default 1.0
+        """
+        self.ksize = ksize
+        self.sigma = sigma
+        self.gamma_1 = 0.8
+        self.gamma_2 = 1.8
+        self.rho = 100.0
+
+    def filter(self, value):
+        return self.gamma_1 - self.gamma_2 * (
+            1 / (1 + 2.415 * np.power(value / self.rho, 4))
+        )
+
+    def _process_image(self, hsi: np.ndarray) -> np.ndarray:
+        """
+        Process image using the model.
+
+        Parameters
+        ----------
+        image : np.ndarray
+            Image to be processed, as HSI.
+
+        Returns
+        -------
+        np.ndarray
+            Processed image, as HSI.
+        """
+        i = hsi[:, :, 2]
+        i_log = np.log2(i + 1.0)
+        i_log_fft_shifted = np.fft.fftshift(np.fft.fft2(i_log))
+        i_log_fft_shifted_filtered = np.zeros_like(i_log_fft_shifted)
+        for i in range(i_log_fft_shifted.shape[0]):
+            for j in range(i_log_fft_shifted.shape[1]):
+                i_log_fft_shifted_filtered[i, j] = i_log_fft_shifted[
+                    i, j
+                ] * self.filter(np.sqrt(i**2 + j**2))
+        i_log_filtered = np.real(
+            np.fft.ifft2(np.fft.ifftshift(i_log_fft_shifted_filtered))
+        )
+        i_filtered = np.exp2(i_log_filtered) - 1.0
+        hsi_filtered = hsi.copy()
+        hsi_filtered[:, :, 2] = i_filtered
+        return hsi_filtered
+
+    def process_image(self, image: np.ndarray) -> np.ndarray:
+        """
+        Process image using the model.
+
+        Parameters
+        ----------
+        image : np.ndarray
+            Image to be processed, as BGR.
+
+        Returns
+        -------
+        np.ndarray
+            Processed image, as BGR.
+        """
+        og = image.copy()
+        image = image.astype(np.float32)
+        hsi = bm.BGR2HSI(image)
+        hsi = self._process_image(hsi)
+        image = bm.HSI2BGR(hsi)
+        image = np.clip(image, 0, 255)
+        image = image.astype(np.uint8)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+        image[:, :, 2] = cv2.equalizeHist(image[:, :, 2])
+        image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
+        # show difference image
+        diff = np.abs(og.astype(np.float32) - image.astype(np.float32))
+        diff = diff.astype(np.uint8)
+        cv2.imshow("diff", diff)
+        cv2.waitKey(0)
+        # show og
+        cv2.imshow("og", og)
+        cv2.waitKey(0)
+        return image

eiuie/main.py

@@ -0,0 +1,65 @@
+import argparse
+
+import cv2
+
+import download as download_module
+import base_model as bm
+import unsharp_masking
+import retinex
+import homomorphic_filtering
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Entry point for the eiuie CLI.")
+
+    parser.add_argument(
+        "command",
+        type=str,
+        choices=["download", "single"],
+        help="Command to run",
+    )
+
+    # --method=xyz
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="unsharp_masking",
+        choices=["unsharp_masking", "retinex", "homomorphic_filtering"],
+        help="Filter method to use",
+    )
+
+    # --file=xyz
+    parser.add_argument(
+        "--file",
+        type=str,
+        help="Path to image file to process",
+    )
+
+    args = parser.parse_args()
+
+    method: bm.BaseModel
+    match args.method:
+        case "unsharp_masking":
+            method = unsharp_masking.UnsharpMasking()
+        case "retinex":
+            method = retinex.Retinex()
+        case "homomorphic_filtering":
+            method = homomorphic_filtering.HomomorphicFiltering()
+        case _:
+            raise ValueError(f"Unknown method: {args.method}")
+
+    match args.command:
+        case "download":
+            download_module.main()
+        case "single":
+            image = cv2.imread(args.file)
+            processed_image = method.process_image(image)
+            # show image
+            cv2.imshow("image", processed_image)
+            cv2.waitKey()
+        case _:
+            raise ValueError(f"Unknown command: {args.command}")
+
+
+if __name__ == "__main__":
+    main()