diff --git a/prettypretty/cloud.py b/prettypretty/cloud.py new file mode 100644 index 0000000..b271d6c --- /dev/null +++ b/prettypretty/cloud.py @@ -0,0 +1,154 @@ +import argparse +import io +from typing import Self + +from prettypretty.color import Color, ColorSpace +from prettypretty.color.gamut import GamutTraversalStep +from prettypretty.color.spectrum import ( + CIE_ILLUMINANT_D65, CIE_OBSERVER_2DEG_1931, SpectrumTraversal +) + + +def create_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + parser.add_argument( + "--gamut", "-g", + help="plot boundary of colorspace gamut; value must be sRGB, P3, or Rec2020", + ) + parser.add_argument( + "--planar", + action="store_true", + help="plot gamut boundary in 2D", + ) + + return parser + + +class PointManager: + def __init__(self, darken: bool = False, planar: bool = False) -> None: + self.darken = darken + self.planar = planar + self.points: list[tuple[float, float, float]] = [] + self.colors: list[tuple[int, int, int]] = [] + self.line_starts: list[int] = [] + + def __len__(self) -> int: + return len(self.points) + + def add(self, color: Color) -> None: + x, y, z = color + + if self.planar: + if color.space().is_ok(): + # Zero out lightness + self.points.append((0, y, z)) + else: + # Zero out luminosity + self.points.append((x, 0, z)) + else: + self.points.append((x, y, z)) + + if self.darken: + lr, c, h = color.to(ColorSpace.Oklrch) + color = Color.oklrch(lr * 2 / 3, c * 2 / 3, h) + + r, g, b = color.to_24bit() + self.colors.append((r, g, b)) + + def mark_newline(self) -> None: + self.line_starts.append(len(self.points)) + + def write_header(self, file: io.TextIOWrapper, include: None | Self = None) -> None: + extra = 0 if include is None else len(include) + + file.write(f"element vertex {len(self.points) + extra}\n") + file.write("property float x\n") + file.write("property float y\n") + file.write("property float z\n") + file.write("property uchar red\n") + file.write("property uchar green\n") + file.write("property uchar blue\n") + + # file.write(f"element face {665}\n") + # file.write("property list uchar int vertex_indices\n") + + def write_data(self, file: io.TextIOWrapper) -> None: + for point, color in zip(self.points, self.colors): + f = 100.0 + x, y, z = point[0] * f, point[1] * f, point[2] * f + r, g, b = color + + file.write(f"{x:f} {y:f} {z:f} {r} {g} {b}\n") + + +def render( + space: ColorSpace, + segment_size: int = 50, + gamut: None | ColorSpace = None, + planar: bool = False, +) -> None: + traversal = SpectrumTraversal(CIE_ILLUMINANT_D65, CIE_OBSERVER_2DEG_1931) + traversal.set_step_sizes(2) + + points = PointManager() + for step in traversal: + if isinstance(step, GamutTraversalStep.MoveTo): + points.mark_newline() + + color = step.color() + if space.is_ok(): + points.add(color.to(ColorSpace.Oklrab)) + else: + points.add(color.to(ColorSpace.Xyz)) + + gamut_points = PointManager(planar=planar) + if gamut is not None: + it = gamut.gamut(segment_size) + assert it is not None + for step in it: + if isinstance(step, GamutTraversalStep.MoveTo): + gamut_points.mark_newline() + + color = step.color() + if space.is_ok(): + gamut_points.add(color.to(ColorSpace.Oklrab)) + else: + gamut_points.add(color.to(ColorSpace.Xyz)) + + if space.is_ok(): + filename = "ok.ply" + else: + filename = "xyz.ply" + + with open(filename, mode="w", encoding="utf8") as file: + file.write("ply\n") + file.write("format ascii 1.0\n") + if gamut is None: + points.write_header(file) + else: + points.write_header(file, gamut_points) + file.write("end_header\n") + + points.write_data(file) + if gamut is not None: + gamut_points.write_data(file) + + +if __name__ == "__main__": + parser = create_parser() + options = parser.parse_args() + if options.gamut is None: + gamut = None + else: + name = options.gamut.lower() + if name == "srgb": + gamut = ColorSpace.Srgb + elif name == "p3": + gamut = ColorSpace.DisplayP3 + elif name == "rec2020": + gamut = ColorSpace.Rec2020 + else: + raise ValueError(f"{name} is not a valid gamut name") + + render(ColorSpace.Xyz, gamut=gamut, planar=options.planar) + render(ColorSpace.Oklrab, gamut=gamut, planar=options.planar)