main.py

#!/usr/bin/env python3
"""
PanoSharp: Unified 3DGS Panorama Alignment Pipeline

A unified command-line interface for aligning and merging 3D Gaussian Splatting
point clouds from Apple's SHARP model. Automatically handles both real photographs
and AI-generated images with appropriate alignment strategies.

Usage Examples:
    # Align real photographs (full pipeline: image -> SHARP -> alignment)
    python main.py --src photo1.jpg --tgt photo2.jpg --type real --output outputs/

    # Align AI-generated images
    python main.py --src ai1.png --tgt ai2.png --type ai --output outputs/

    # Use pre-existing PLY files (skip SHARP generation)
    python main.py --src_ply outputs/1.ply --tgt_ply outputs/2.ply \\
                   --src_img img1.jpg --tgt_img img2.jpg \\
                   --type real --output outputs/

    # Customize alignment parameters
    python main.py --src img1.jpg --tgt img2.jpg --type real --output outputs/ \\
                   --smoothing 1.0 --warp-strength 0.5 --falloff 0.5

Alignment Strategies:
    --type real: For photographs from cameras
        - Uses Sim3 transform (rotation + translation + scale)
        - Depth affine correction for monocular depth bias
        - Non-rigid RBF warping for residual correction
    
    --type ai: For AI-generated images (DALL-E, Midjourney, Gemini, etc.)
        - Uses 2D homography (robust to depth inconsistencies)
        - Scale estimation from depth ratios
        - 2D RBF smoothing

Output Files:
    {output_dir}/
    ├── ply/
    │   ├── {src_name}.ply    # SHARP output for source image
    │   └── {tgt_name}.ply    # SHARP output for target image
    ├── aligned.ply           # Source PLY aligned to target
    └── merged.ply            # Combined target + aligned source

Author: PanoSharp Team
License: See LICENSE file
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch

from panosharp.core.data_types import ImageType, PipelineInput
from panosharp.pipeline.orchestrator import PanoSharpPipeline


def parse_args() -> argparse.Namespace:
    """
    Parse command-line arguments.

    Returns:
        Parsed arguments namespace
    """
    parser = argparse.ArgumentParser(
        description="PanoSharp: Unified 3DGS Panorama Alignment Pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Full pipeline with real photos
  python main.py --src photo1.jpg --tgt photo2.jpg --type real --output outputs/

  # AI-generated images with custom parameters
  python main.py --src ai1.png --tgt ai2.png --type ai --output outputs/ \\
                 --smoothing 1.0 --warp-strength 0.5

  # Use existing PLY files (skip SHARP)
  python main.py --src_ply src.ply --tgt_ply tgt.ply \\
                 --src_img src.jpg --tgt_img tgt.jpg \\
                 --type real --output outputs/
        """,
    )

    # Input images (required for full pipeline)
    input_group = parser.add_argument_group("Input Images")
    input_group.add_argument(
        "--src",
        "--src_img",
        dest="src_img",
        type=Path,
        required=True,
        help="Source image path (to be aligned)",
    )
    input_group.add_argument(
        "--tgt",
        "--tgt_img",
        dest="tgt_img",
        type=Path,
        required=True,
        help="Target image path (reference)",
    )

    # Pre-existing PLY files (optional, skip SHARP if provided)
    ply_group = parser.add_argument_group("PLY Files (Optional)")
    ply_group.add_argument(
        "--src_ply",
        type=Path,
        default=None,
        help="Pre-existing source PLY file (skip SHARP generation)",
    )
    ply_group.add_argument(
        "--tgt_ply",
        type=Path,
        default=None,
        help="Pre-existing target PLY file (skip SHARP generation)",
    )

    # Output
    output_group = parser.add_argument_group("Output")
    output_group.add_argument(
        "--output",
        "-o",
        type=Path,
        required=True,
        help="Output directory for results",
    )

    # Image type (required)
    type_group = parser.add_argument_group("Image Type")
    type_group.add_argument(
        "--type",
        "-t",
        choices=["real", "ai"],
        required=True,
        help="Image type: 'real' for photographs, 'ai' for AI-generated",
    )

    # Device
    device_group = parser.add_argument_group("Computation")
    device_group.add_argument(
        "--device",
        choices=["cuda", "cpu", "mps"],
        default="cuda",
        help="Computation device (default: cuda)",
    )

    # Alignment parameters
    align_group = parser.add_argument_group("Alignment Parameters")
    align_group.add_argument(
        "--smoothing",
        type=float,
        default=0.3,
        help="RBF smoothing (higher = smoother warp, default: 0.3)",
    )
    align_group.add_argument(
        "--warp-strength",
        type=float,
        default=1.0,
        help="Warp intensity (0 = rigid only, default: 1.0)",
    )
    align_group.add_argument(
        "--falloff",
        type=float,
        default=0.5,
        help="Distance falloff (0 = no falloff, default: 0.0)",
    )
    align_group.add_argument(
        "--scale",
        type=float,
        default=0.0,
        help="Manual depth scale for AI images (0 = auto, default: 0)",
    )

    # SHARP options
    sharp_group = parser.add_argument_group("SHARP Options")
    sharp_group.add_argument(
        "--sharp-checkpoint",
        type=Path,
        default=None,
        help="Path to SHARP model checkpoint (optional)",
    )

    return parser.parse_args()


def validate_args(args: argparse.Namespace) -> None:
    """
    Validate parsed arguments.

    Args:
        args: Parsed arguments

    Raises:
        SystemExit: If validation fails
    """
    # Check source image exists
    if not args.src_img.exists():
        print(f"Error: Source image not found: {args.src_img}", file=sys.stderr)
        sys.exit(1)

    # Check target image exists
    if not args.tgt_img.exists():
        print(f"Error: Target image not found: {args.tgt_img}", file=sys.stderr)
        sys.exit(1)

    # If PLY files provided, check they exist
    if args.src_ply and not args.src_ply.exists():
        print(f"Error: Source PLY not found: {args.src_ply}", file=sys.stderr)
        sys.exit(1)

    if args.tgt_ply and not args.tgt_ply.exists():
        print(f"Error: Target PLY not found: {args.tgt_ply}", file=sys.stderr)
        sys.exit(1)

    # Must provide both PLY files or neither
    if (args.src_ply is None) != (args.tgt_ply is None):
        print(
            "Error: Must provide both --src_ply and --tgt_ply or neither",
            file=sys.stderr,
        )
        sys.exit(1)

    # Check device availability
    if args.device == "cuda" and not torch.cuda.is_available():
        print("Warning: CUDA not available, falling back to CPU", file=sys.stderr)
        args.device = "cpu"

    if args.device == "mps" and not torch.backends.mps.is_available():
        print("Warning: MPS not available, falling back to CPU", file=sys.stderr)
        args.device = "cpu"


def main() -> int:
    """
    Main entry point for PanoSharp CLI.

    Returns:
        Exit code (0 = success, 1 = failure)
    """
    # Parse and validate arguments
    args = parse_args()
    validate_args(args)

    # Convert image type string to enum
    image_type = ImageType.REAL if args.type == "real" else ImageType.AI_GENERATED

    # Create pipeline input
    pipeline_input = PipelineInput(
        src_image_path=args.src_img,
        tgt_image_path=args.tgt_img,
        output_dir=args.output,
        image_type=image_type,
        src_ply_path=args.src_ply,
        tgt_ply_path=args.tgt_ply,
        device=args.device,
        smoothing=args.smoothing,
        warp_strength=args.warp_strength,
        falloff=args.falloff,
        scale_factor=args.scale,
    )

    # Create and run pipeline
    pipeline = PanoSharpPipeline.create_default(
        device=args.device,
        sharp_checkpoint=args.sharp_checkpoint,
    )

    result = pipeline.run(pipeline_input)

    # Print summary
    print("\n" + "=" * 60)
    if result.success:
        print("SUCCESS!")
        print("=" * 60)
        print(f"Source PLY: {result.src_ply_path}")
        print(f"Target PLY: {result.tgt_ply_path}")
        print(f"Aligned: {result.aligned_ply_path}")
        if result.merged_ply_path:
            print(f"Merged: {result.merged_ply_path}")

        stats = result.alignment_result.statistics
        print(f"\nAlignment Statistics:")
        print(f"  Matches: {stats.num_matches}")
        print(f"  Inliers: {stats.num_inliers}")
        print(f"  Transform: {stats.transform_type}")
        print(f"  Residual: mean={stats.residual_mean:.4f}, max={stats.residual_max:.4f}")
        print(f"  Warp: mean={stats.warp_mean:.4f}, max={stats.warp_max:.4f}")

        if stats.depth_affine:
            print(f"  Depth affine: a={stats.depth_affine[0]:.4f}, b={stats.depth_affine[1]:.4f}")

        return 0
    else:
        print("FAILED!")
        print("=" * 60)
        print(f"Error: {result.error_message}")
        return 1


if __name__ == "__main__":
    sys.exit(main())