Bundle adjustment in Julia (doesn't work yet!)

bundle_adjustment.jl

#
#  Bundle adjustment
#

using JuMP
using StaticArrays
using Rotations
using LinearAlgebra
import Ipopt

"Euclidean coordinates for a point in R^2."
Point2 = SVector{2, Float64}

"Euclidean coordinates for a point in R^3."
Point3 = SVector{3, Float64}

abstract type Camera end

# "Parameters for a pinhole camera."
mutable struct PinholeCamera <: Camera
    focalLength :: Float64
    principalPoint :: Tuple{Float64, Float64}
    location :: Point3  # coordinates of camera center in the world coordinate frame
    orientation :: RodriguesVec{Float64}

    PinholeCamera() =
        new(1.0,
            (0.0, 0.0),
            Point3(0, 0, 0),
            one(RotMatrix{3, Float64}))
end

calibrationmatrix(cam::PinholeCamera) =
    SMatrix{3,3}([
      cam.focalLength 0.0             cam.principalPoint[1];
      0.0             cam.focalLength cam.principalPoint[2];
      0.0             0.0             1.0;
    ])

cameramatrix(cam::PinholeCamera) =
    calibrationmatrix(cam) * cam.orientation * hcat(I, -cam.location)

mutable struct Bundle{C <: Camera}
    cameras :: Array{C, 1}    # Cameras' locations in world coordinate space
    world_points :: Array{Float64, 2}   # 3D points in world coordinate space
end

function initialize_bundle(::Type{PinholeCamera}, matches::AbstractArray{Float64, 3})
    npoints, ncameras, ncoords = size(matches)
    @assert ncoords == 2 "`matches` must have shape (# points, # cameras, 2), but the 3rd axis has size $ncoords"

    cameras = [PinholeCamera() for _ in 1 : ncameras]

    # TODO Backproject the features with a fixed distance to the camera center,
    # rather than just setting their Z coordinate
    worldpoints = zeros(npoints, ncoords + 1)
    worldpoints[:, 1:2] = matches[:, 1, :]
    worldpoints[:, 3] .= 5.0  # set a default depth

    return Bundle{PinholeCamera}(cameras, worldpoints)
end

function create_model(bundle::Bundle{PinholeCamera},
                      matches::AbstractArray{Float64, 3})
    npoints, ncameras, ncoords = size(matches)
    @assert ncoords == 2 "`matches` must have shape (# points, # cameras, 2), but the 3rd axis has size $ncoords"
    @assert ncameras == length(bundle.cameras)  "Number of cameras differ between `matches` and `bundle.cameras`"

    # The problem is, roughly:
    #
    # minimize camera[*].location, camera[*].orientation, map.point[*]
    #   sum i=1:n_cameras
    #     sum j=1:n_keypoints
    #       d^2(camera[i].feature[j], project(camera[i], map.point[j]))
    #
    # where camera[i].location    :: [3] = i-th camera's location
    #       camera[i].orientation :: [3] = i-th camera's orientation (Euler angles)
    #       map.point[i]          :: [4] = i-th 3D point's coordinates in world
    #                                      coordinate frame, homogeneous coordinates

    model = Model(with_optimizer(Ipopt.Optimizer))

    # Camera locations (Eulerian coordinates)
    camloc = @variable(model, camloc[1:ncameras, 1:3])
    # Camera orientations (3×3 rotation matrices)
    camrot = @variable(model, camrot[1:ncameras, 1:3])
    # Camera focal lengths (scalar)
    camfl = @variable(model, camfl[1:ncameras])
    # Camera principal points (2-tuple, (x, y))
    campp = @variable(model, campp[1:ncameras, 1:2])

    for i = 1 : ncameras
        if i == 1
            @NLconstraint(model, [k = 1:3], camloc[i, k] == bundle.cameras[i].location[k])
            @NLconstraint(model, [k = 1:3], camrot[i, k] == bundle.cameras[i].orientation[k])
        else
            for k = 1:3
                set_start_value(camloc[i, k], bundle.cameras[i].location[k])
                set_start_value(camrot[i, k], bundle.cameras[i].orientation[k])
            end
        end

        # TODO focalLength should have both (x, y) coordinates
        set_start_value(camfl[i], bundle.cameras[i].focalLength)
        for k = 1:2
            set_start_value(campp[i, k], bundle.cameras[i].principalPoint[k])
        end
    end

    wpoints = @variable(model, wpoints[1:npoints, 1:4])
    for i = 1:npoints
        for j = 1:3
            set_start_value(wpoints[i, j], bundle.world_points[i, j])
        end
        set_start_value(wpoints[i, 4], 1.0)
    end

    K(focal_length, principal_point_x, principal_point_y) =
        SMatrix{3,3}([
          focal_length 0.0          principal_point_x;
          0.0          focal_length principal_point_y;
          0.0          0.0          1.0;
        ])
    register(model, :K, 2, K, autodiff=true)

    project(fl, ppx, ppy, loc, rot, point) =
        K(fl, ppx, ppy) * rot * (point - loc)
    register(model, :project, 6, project, autodiff=true)

    distance_sq(x, y) = sum((x - y) .^ 2)
    register(model, :distance_sq, 2, distance_sq, autodiff=true)
    
    @variable(model, cost)
    @NLconstraint(model, cost ==
        sum(distance_sq(
                matches[j, i, :],
                project(
                    camfl[i],
                    campp[i,1], campp[i,2],
                    camloc[i,:],
                    camrot[i,:],
                    wpoints[j,:]))
            for i = 1:ncameras
            for j = 1:npoints))
    @NLobjective(model, Min, cost)

    model
end