WUT_Computer_Science/Programming/TRAK/main.py

import argparse
import importlib
import os

import matplotlib.pyplot as plt
from photon_mapping import render_photon_mapping
from rendering import ray_trace
from utils import load_config, parse_resolution


def main():
    # default config
    config = load_config("config.ini")

    # Parse
    parser = argparse.ArgumentParser(description="Rendering Program")
    parser.add_argument(
        "--algorithm",
        type=str,
        help="Algorithm to use",
        default=config.get("DEFAULT", "algorithm"),
    )
    parser.add_argument(
        "--scene",
        type=str,
        help="Name of the scene to render (without .py).",
        default=config.get("DEFAULT", "scene"),
    )
    parser.add_argument(
        "--environment",
        type=str,
        help="Environment file",
        default=config.get("DEFAULT", "environment"),
    )
    parser.add_argument(
        "--env_blur",
        type=str,
        help="Environment blur",
        default=config.get("DEFAULT", "env_blur"),
    )
    parser.add_argument(
        "--resolution",
        type=str,
        help="Image resolution (WIDTHxHEIGHT)",
        default=config.get("DEFAULT", "resolution"),
    )
    parser.add_argument(
        "--samples_per_pixel",
        type=int,
        default=config.get("ray_tracing", "samples_per_pixel"),
        help="Samples per pixel for rendering.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=config.get("DEFAULT", "output"),
        help="Output file name.",
    )

    parser.add_argument(
        "--num_spheres",
        type=int,
        default=3,
        help="Number of spheres in the scene for Ray Tracing 0",
    )
    parser.add_argument(
        "--num_photons",
        type=int,
        default=config.getint("photon_mapping", "num_photons"),
        help="Number of photons for photon mapping",
    )
    parser.add_argument(
        "--max_depth",
        type=int,
        default=config.getint("photon_mapping", "max_depth"),
        help="Maximum depth for photon tracing",
    )
    parser.add_argument(
        "--gather_radius",
        type=float,
        default=config.getfloat("photon_mapping", "gather_radius"),
        help="Radius for radiance estimation in photon mapping",
    )
    # The render is always written to outputs/; this only controls whether a
    # viewer is ALSO popped up. It defaults off because img.show() shells out
    # to xdg-open and then blocks until a human closes the window -- which
    # left ../../run.sh run Programming/TRAK hanging for 45 minutes after the
    # ray trace had actually finished.
    parser.add_argument(
        "--show",
        action="store_true",
        help="also open the finished image in a viewer (blocks until closed)",
    )

    args = parser.parse_args()

    width, height = parse_resolution(args.resolution)

    # Run the selected algorithm
    if args.algorithm == "ray_tracing0":
        print("Starting ray tracing zero...")
        # ray_trace(args.scene, args.environment_map, image_width=width, image_height=height, output_file="output_ray_traced.png")
        ray_trace(
            args.num_spheres,
            args.environment_map,
            image_width=width,
            image_height=height,  # na razie generujemy w kodzie, ale potem trzeba będzie obj wczytywać
            output_file="output_ray_traced.png",
        )
    elif args.algorithm == "ray_tracing":
        print("Starting ray tracing...")
        try:
            print(args.scene)
            scene_module = importlib.import_module(f"scenes.{args.scene}")
        except ModuleNotFoundError:
            print(f"Error: Scene '{args.scene}' not found in the 'scenes' directory.")
            return
        try:
            scene = scene_module.setup_scene(
                width=width, height=height, environment=f"{args.environment}"
            )
        except AttributeError:
            print(
                f"Error: Scene '{args.scene}' does not define a `setup_scene` function."
            )
            return
        # Renderowanie
        print(
            f"Rendering scene '{args.scene}' with {args.samples_per_pixel} samples per pixel..."
        )
        img = scene.render(samples_per_pixel=args.samples_per_pixel)
        output_path = os.path.join("outputs", args.output)
        img.save(output_path)
        print(f"Image saved to {output_path}")
        if args.show:
            img.show()
    elif args.algorithm == "photon_mapping":
        print("Starting photon mapping...")
        image = render_photon_mapping(
            width, height, args.num_photons, args.max_depth, args.gather_radius
        )
        plt.imshow(image)
        plt.axis("off")
        output_path = os.path.join("outputs", args.output)
        plt.savefig(output_path)
        print(f"Image saved to {output_path}")
        if args.show:
            plt.show()
    else:
        print(f"Unknown algorithm: {args.algorithm}")
        return


if __name__ == "__main__":
    main()