projection mapping

#!/usr/bin/env python3
"""
projection_prewarp.py

Interactive tool to compute a homography between a source image and a photo of a projected result,
then create a pre-warped image so that after projector-perspective the result appears correct.

Features:
 - Manual or automatic calibration (point correspondences)
 - Save and load homography
 - Batch mode: apply existing homography to new images
 - Interactive viewer with zoom and pan for precise point selection
 - Undo function to remove last point

Usage examples:
 Manual calibration:
   python projection_prewarp.py --source source.png --photo projected.jpg --output prewarped.png

 Automatic feature matching:
   python projection_prewarp.py --source source.png --photo projected.jpg --output prewarped.png --mode auto

 Apply precomputed homography:
   python projection_prewarp.py --apply result_H.npy --input new.png --output new_prewarped.png --outsize 1920x1080

"""

import argparse
import json
import os
import sys
from datetime import datetime

import cv2
import numpy as np

# --- Globals used by mouse callbacks ---
src_points = []
dst_points = []

class ImageViewer:
    def __init__(self, name, img, points_list):
        self.name = name
        self.img = img
        self.points_list = points_list
        self.zoom = 1.0
        self.offset = np.array([0, 0], dtype=np.float32)
        self.dragging = False
        self.drag_start = None
        cv2.namedWindow(self.name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.name, self.on_mouse)

    def on_mouse(self, event, x, y, flags, param=None):
        if event == cv2.EVENT_LBUTTONDOWN:
            orig_x = int((x + self.offset[0]) / self.zoom)
            orig_y = int((y + self.offset[1]) / self.zoom)
            self.points_list.append((orig_x, orig_y))
        elif event == cv2.EVENT_RBUTTONDOWN:
            self.dragging = True
            self.drag_start = np.array([x, y], dtype=np.float32)
        elif event == cv2.EVENT_MOUSEMOVE and self.dragging:
            delta = np.array([x, y], dtype=np.float32) - self.drag_start
            self.offset -= delta
            self.offset = np.maximum(self.offset, 0)
            self.drag_start = np.array([x, y], dtype=np.float32)
        elif event == cv2.EVENT_RBUTTONUP:
            self.dragging = False
        elif event == cv2.EVENT_MOUSEWHEEL:
            if flags > 0:
                self.zoom *= 1.2
            else:
                self.zoom /= 1.2
            self.zoom = max(0.2, min(self.zoom, 10.0))

    def show(self):
        h, w = self.img.shape[:2]
        view = cv2.resize(self.img, (int(w * self.zoom), int(h * self.zoom)))
        # Crop window to max 1200x800
        x0 = int(self.offset[0])
        y0 = int(self.offset[1])
        x1 = x0 + min(view.shape[1], 1200)
        y1 = y0 + min(view.shape[0], 800)
        crop = view[y0:y1, x0:x1].copy()
        for i, (px, py) in enumerate(self.points_list):
            sx = int(px * self.zoom - self.offset[0])
            sy = int(py * self.zoom - self.offset[1])
            if 0 <= sx < crop.shape[1] and 0 <= sy < crop.shape[0]:
                cv2.circle(crop, (sx, sy), 5, (0, 255, 0), -1)
                cv2.putText(crop, str(i + 1), (sx + 8, sy - 8),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 1)
        cv2.imshow(self.name, crop)

def save_points(outfile, src_pts, dst_pts):
    data = {
        'timestamp': datetime.utcnow().isoformat() + 'Z',
        'src_points': [[float(x), float(y)] for (x, y) in src_pts],
        'dst_points': [[float(x), float(y)] for (x, y) in dst_pts],
    }
    with open(outfile, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2)
    print(f"Saved points to {outfile}")

def try_auto_match(src_img_gray, dst_img_gray, max_features=2000, ratio_test=True):
    orb = cv2.ORB_create(nfeatures=max_features)
    kp1, des1 = orb.detectAndCompute(src_img_gray, None)
    kp2, des2 = orb.detectAndCompute(dst_img_gray, None)
    print(f"ORB: found {len(kp1)} keypoints in source, {len(kp2)} in photo")
    if des1 is None or des2 is None or len(kp1) < 4 or len(kp2) < 4:
        print("Not enough keypoints/descriptors for automatic matching")
        return None
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
    try:
        matches = bf.knnMatch(des1, des2, k=2)
    except Exception:
        matches = [[m] for m in bf.match(des1, des2)]
    good = []
    if ratio_test:
        for m_n in matches:
            if len(m_n) >= 2:
                m, n = m_n[0], m_n[1]
                if m.distance < 0.75 * n.distance:
                    good.append(m)
    else:
        for m_n in matches:
            good.append(m_n[0])
    print(f"Auto-matcher found {len(good)} good matches")
    if len(good) < 8:
        return None
    pts_src = np.float32([kp1[m.queryIdx].pt for m in good])
    pts_dst = np.float32([kp2[m.trainIdx].pt for m in good])
    return pts_src, pts_dst, kp1, kp2, good

def compute_and_warp(src, photo, pts_src, pts_dst, out_size=None, save_prefix='result'):
    pts_src = np.array(pts_src, dtype=np.float32)
    pts_dst = np.array(pts_dst, dtype=np.float32)
    if pts_src.shape[0] < 4 or pts_dst.shape[0] < 4:
        raise ValueError('Need at least 4 point correspondences')
    H, mask = cv2.findHomography(pts_src, pts_dst, cv2.RANSAC, 5.0)
    if H is None:
        raise RuntimeError('findHomography failed')
    print('Homography (src -> photo):')
    print(H)
    np.save(save_prefix + '_H.npy', H)
    with open(save_prefix + '_H.txt', 'w') as f:
        f.write('\n'.join([' '.join(map(str, row)) for row in H]))
    try:
        H_inv = np.linalg.inv(H)
    except np.linalg.LinAlgError:
        ok, H_inv = cv2.invert(H)
        if not ok:
            raise RuntimeError('Homography is singular')
    if out_size is None:
        h_out, w_out = photo.shape[:2]
    else:
        w_out, h_out = out_size
    prewarped = cv2.warpPerspective(src, H_inv, (w_out, h_out), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT)
    cv2.imwrite(save_prefix + '_prewarped.png', prewarped)
    return H, H_inv, prewarped

def apply_existing_homography(h_file, input_file, output_file, out_size=None):
    H = np.load(h_file)
    try:
        H_inv = np.linalg.inv(H)
    except np.linalg.LinAlgError:
        ok, H_inv = cv2.invert(H)
        if not ok:
            raise RuntimeError('Homography is singular')
    img = cv2.imread(input_file)
    if img is None:
        raise RuntimeError(f'Failed to load {input_file}')
    if out_size is None:
        out_w, out_h = img.shape[1], img.shape[0]
    else:
        out_w, out_h = out_size
    prewarped = cv2.warpPerspective(img, H_inv, (out_w, out_h), flags=cv2.INTER_LINEAR)
    cv2.imwrite(output_file, prewarped)
    print(f"Applied homography from {h_file} to {input_file}, saved {output_file}")

def main():
    parser = argparse.ArgumentParser(description='Compute or apply homography for projector prewarp')
    parser.add_argument('--source', '-s', help='Source image (for calibration)')
    parser.add_argument('--photo', '-p', help='Photo of projected image (for calibration)')
    parser.add_argument('--output', '-o', default='prewarped.png', help='Output filename for prewarped image')
    parser.add_argument('--mode', choices=['manual', 'auto'], default='manual', help='Calibration mode')
    parser.add_argument('--outsize', help='Output size WIDTHxHEIGHT for prewarped image')
    parser.add_argument('--save-prefix', default='result', help='Prefix for saved calibration files')
    parser.add_argument('--apply', help='Use existing homography file (.npy)')
    parser.add_argument('--input', help='Input image to prewarp using existing homography')
    args = parser.parse_args()

    out_size = None
    if args.outsize:
        try:
            w, h = args.outsize.split('x')
            out_size = (int(w), int(h))
        except Exception:
            print('Invalid --outsize, expected WIDTHxHEIGHT')
            sys.exit(1)

    if args.apply and args.input:
        apply_existing_homography(args.apply, args.input, args.output, out_size)
        return

    if not args.source or not args.photo:
        print('For calibration you must provide --source and --photo')
        sys.exit(1)

    src_img = cv2.imread(args.source, cv2.IMREAD_COLOR)
    dst_img = cv2.imread(args.photo, cv2.IMREAD_COLOR)
    if src_img is None or dst_img is None:
        print('Failed to load images')
        sys.exit(1)

    if args.mode == 'auto':
        gray1 = cv2.cvtColor(src_img, cv2.COLOR_BGR2GRAY)
        gray2 = cv2.cvtColor(dst_img, cv2.COLOR_BGR2GRAY)
        auto = try_auto_match(gray1, gray2)
        if auto:
            pts_src, pts_dst, kp1, kp2, good_matches = auto
            matches_img = cv2.drawMatches(src_img, kp1, dst_img, kp2, good_matches[:60], None,
                                          flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
            cv2.imshow('Matches (auto)', matches_img)
            key = cv2.waitKey(0)
            cv2.destroyWindow('Matches (auto)')
            if key != 27:
                H, H_inv, prewarped = compute_and_warp(src_img, dst_img, pts_src, pts_dst, out_size, args.save_prefix)
                cv2.imwrite(args.output, prewarped)
                save_points(args.save_prefix + '_points.json', pts_src.tolist(), pts_dst.tolist())
                return
        print('Auto failed or cancelled, switching to manual')

    src_view = ImageViewer('Source', src_img, src_points)
    dst_view = ImageViewer('Photo', dst_img, dst_points)

    while True:
        src_view.show()
        dst_view.show()
        key = cv2.waitKey(20) & 0xFF
        if key == ord('r'):
            src_points.clear()
            dst_points.clear()
        elif key == ord('s'):
            save_points(args.save_prefix + '_points.json', src_points, dst_points)
        elif key == ord('c'):
            if len(src_points) >= 4 and len(src_points) == len(dst_points):
                H, H_inv, prewarped = compute_and_warp(src_img, dst_img, src_points, dst_points, out_size, args.save_prefix)
                cv2.imwrite(args.output, prewarped)
                save_points(args.save_prefix + '_points.json', src_points, dst_points)
            else:
                print('Need >=4 pairs and equal count')
            break
        elif key == ord('u') or key == 8:  # Undo last point (u or Backspace)
            if src_points and (len(src_points) > len(dst_points)):
                removed = src_points.pop()
                print(f"Cofnięto punkt źródłowy {removed}")
            elif dst_points:
                removed = dst_points.pop()
                print(f"Cofnięto punkt docelowy {removed}")
        elif key == ord('q') or key == 27:
            break
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()