CLAHE image post-processing plugin for Pupil Player

README.md

CLAHE image post-processing plugin for Pupil Player

Contrast Limited Adaptive Histogram Equalization - Read more about Histogram Equalization here.

Preview plugin

File: CLAHE-correction-plugin.py

This plugin applies the CLAHE algorithm to the scene video previewed in Pupil Player. See this documentation on how to install the plugin.

Application script

File: apply-clahe.py

This Player-independent script reads the scene video, applies the algorithm with a specific parameter set, and writes the result to a new video file. By following these steps, you can replace the original scene video with the post-processed one:

1. Install the requirements (only required once)

• pip install av opencv-python rich

2. Stop Pupil Player (in case it is running)
3. Rename world.mp4 to backup.mp4
4. Run the script

• python apply-clahe.py <path to recording folder> -cl 40 -gsr 8 -gsc 8
• The result will be written to world.mp4 in the given recording folder
• Use the cl, gsr, gsc options to set custom CLAHE parameters

5. Delete the world_lookup.npy file
6. Open the recording in Pupil Player

General usage:

usage: apply-clahe.py [-h] [-cl CLIP_LIMIT] [-gsr TILE_GRID_SIZE_ROW] [-gsc TILE_GRID_SIZE_COL] [--in-file IN_FILE] [--out-file OUT_FILE] [--force]
                      recording

positional arguments:
  recording

optional arguments:
  -h, --help            show this help message and exit
  -cl CLIP_LIMIT, --clip-limit CLIP_LIMIT
  -gsr TILE_GRID_SIZE_ROW, --tile-grid-size-row TILE_GRID_SIZE_ROW
  -gsc TILE_GRID_SIZE_COL, --tile-grid-size-col TILE_GRID_SIZE_COL
  --in-file IN_FILE
  --out-file OUT_FILE
  --force

apply-clahe.py

import argparse
import pathlib
from struct import pack

import av
import cv2
from rich.progress import track


def main(
    recording,
    in_file,
    out_file,
    force=False,
    clip_limit=40.0,
    tile_grid_size_row=8,
    tile_grid_size_col=8,
):
    recording = pathlib.Path(recording).resolve()
    in_file = recording / in_file
    out_file = recording / out_file
    if not in_file.exists():
        raise FileNotFoundError(f"Input file does not exist: {in_file}")
    if out_file.exists() and not force:
        raise FileExistsError(f"Output file exists already. Use --force to overwrite.")

    with av.open(str(in_file)) as in_container:
        with av.open(str(out_file), "w") as out_container:
            in_stream = in_container.streams.video[0]
            out_stream = out_container.add_stream("mjpeg", rate=in_stream.rate)
            out_stream.time_base = in_stream.time_base
            out_stream.format = in_stream.format

            clahe = cv2.createCLAHE(
                clipLimit=clip_limit,
                tileGridSize=(tile_grid_size_row, tile_grid_size_col),
            )

            for in_frame in track(
                in_container.decode(video=0),
                description="Processing...",
                total=in_container.streams.video[0].frames,
            ):
                gray = in_frame.to_ndarray(format="gray")
                gray = clahe.apply(gray)

                out_frame = av.VideoFrame.from_ndarray(gray, format="gray")
                out_frame.time_base = in_frame.time_base
                out_frame.pts = in_frame.pts
                out_frame = out_frame.reformat(format=out_stream.format)
                for packet in out_stream.encode(out_frame):
                    packet.pts = in_frame.pts
                    out_container.mux(packet)

            for packet in out_stream.encode(None):
                out_container.mux(packet)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("recording")
    parser.add_argument("-cl", "--clip-limit", default=40, type=float)
    parser.add_argument("-gsr", "--tile-grid-size-row", default=8, type=int)
    parser.add_argument("-gsc", "--tile-grid-size-col", default=8, type=int)
    parser.add_argument("--in-file", default="backup.mp4")
    parser.add_argument("--out-file", default="world.mp4")
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    main(
        args.recording,
        args.in_file,
        args.out_file,
        args.force,
        args.clip_limit,
        args.tile_grid_size_row,
        args.tile_grid_size_col,
    )

CLAHE-correction-plugin.py

import cv2
import numpy as np

from plugin import Plugin
from pyglui import ui
from pupil_apriltags import Detector


class CLAHE_Preview(Plugin):
    order = 0.1

    def __init__(
        self,
        g_pool,
        clip_limit: float = 40.0,
        tile_grid_size_row: int = 8,
        tile_grid_size_col: int = 8,
    ):
        super().__init__(g_pool)
        self.clip_limit = clip_limit
        self.tile_grid_size_row = tile_grid_size_row
        self.tile_grid_size_col = tile_grid_size_col
        self.enabled = True
        self.at_detector = Detector(
            families="tag36h11",
            nthreads=1,
            quad_decimate=1.0,
            quad_sigma=0.0,
            refine_edges=1,
            decode_sharpening=0.25,
            debug=0,
        )

    def init_ui(self):
        self.add_menu()
        self.menu.label = "CLAHE Post-Processing Preview"
        self.menu.append(
            ui.Info_Text("Contrast Limited Adaptive Histogram Equalization")
        )
        self.menu.append(
            ui.Info_Text(
                "Read more about Histogram Equalization here: "
                "https://docs.opencv.org/4.x/d5/daf/tutorial_py_histogram_equalization.html"
            )
        )
        self.menu.append(ui.Switch("enabled", self, label="Enable CLAHE"))
        self.menu.append(ui.Text_Input("clip_limit", self, label="Clip Limit"))
        self.menu.append(
            ui.Text_Input("tile_grid_size_row", self, label="Tile Grid Size (rows)")
        )
        self.menu.append(
            ui.Text_Input("tile_grid_size_col", self, label="Tile Grid Size (columns)")
        )

    def deinit_ui(self):
        self.remove_menu()

    def recent_events(self, events):
        frame = events.get("frame")
        if not frame:
            return

        gray = cv2.cvtColor(frame.img, cv2.COLOR_BGR2GRAY)

        if self.enabled:
            clahe = cv2.createCLAHE(
                clipLimit=self.clip_limit,
                tileGridSize=(self.tile_grid_size_row, self.tile_grid_size_col),
            )

            gray = clahe.apply(gray)

        detections = self.at_detector.detect(gray)
        gray = np.repeat(gray[..., np.newaxis], 3, axis=2)

        for detection in detections:
            cv2.polylines(
                gray,
                [detection.corners.reshape((-1, 1, 2)).astype("int32")],
                True,
                (0, 0, 255),
            )

        frame.img[:] = gray

    def get_init_dict(self):
        return {
            "clip_limit": self.clip_limit,
            "tile_grid_size_row": self.tile_grid_size_row,
            "tile_grid_size_col": self.tile_grid_size_col,
        }