comic.jl

comic.jl

using Colors
using CoordinateTransformations
using Images
using LinearAlgebra
using PolygonAlgorithms
using SparseArrays
using StaticArrays

include("linear_rgb.jl")
using .LinearRGB

const INPUT_IMAGE_PATH::String = "$(@__DIR__)/../images/Comic 2 Alpha.png"
const OUTPUT_IMAGE_PATH::String = "$(@__DIR__)/../images/Comic 3 Recursion.png"

# The coordinate transformation.
const FROM_COORDS::SVector{2,NTuple{2,Float64}} = SA[(731, 97), (668, 318)]
const TO_COORDS::SVector{2,NTuple{2,Float64}} = SA[(792, 1), (792, 540)]

# A brightness/contrast effect.
const COLOR_MULTIPLY::RGB{Float64} = RGB{Float64}(0.8, 0.8, 0.8)
const COLOR_ADD::RGB{Float64} = RGB{Float64}(0.02, 0.02, 0.02)

@kwdef struct WeightedPixel
  vector_ix::Int64
  weight::Float64
end

@kwdef struct PixelEquation
  left_side::Vector{WeightedPixel}
  right_side::Float64
end

const SQRT_EPS::Float64 = sqrt(eps())

# The index of the red color component of the pixel coordinate (y, x) within
# the vector used by the solver. The indices of the green and blue color
# components are (index + 1) and (index + 2) respectively.
vector_ix(width::Int64, y::Int64, x::Int64) = ((y - 1) * width + (x - 1)) * 3 + 1

function comic()
  println("Loading $INPUT_IMAGE_PATH")
  img_srgb = load(INPUT_IMAGE_PATH)
  img = srgb_to_linear_rgb.(RGBA{Float64}, img_srgb)

  result_img = process(img)

  println("Saving $OUTPUT_IMAGE_PATH")
  result_img_srgb = linear_rgb_to_srgb.(RGB{N0f8}, result_img)
  save(OUTPUT_IMAGE_PATH, result_img_srgb)

  println("Optimizing $OUTPUT_IMAGE_PATH")
  run(`optipng $OUTPUT_IMAGE_PATH`)
end

function process(img::Matrix{RGBA{Float64}})
  transform = kabsch(SVector.(FROM_COORDS) => SVector.(TO_COORDS), scale=true)

  solve_system(img, transform)
end

# Solve a system of linear equations which equates the result pixels to the
# weighted average of the corresponding pixels after an affine transformation,
# with the input image alpha blended on top and a brightness/contrast effect
# applied according to COLOR_MULTIPLY and COLOR_ADD.
#
# In pseudocode,
#
# dst[y, x] =
#   blend(
#     src[y, x].alpha,
#     sum(area * dst[y_tr, x_tr]) * multiply + add,
#     src[y, x].color
#    )
#
# where (y_tr, x_tr) is each transformed pixel coordinate and area is the
# fraction of that pixel's contribution within the (y, x) pixel.

function solve_system(
  img::Matrix{RGBA{Float64}},
  transform::AbstractAffineMap,
)
  img_height, img_width = size(img)

  println("Constructing system of equations")
  matrix, right_side = @time construct_system(img, transform)

  println("Solving system of equations")
  solution = @time lu(matrix) \ right_side

  result_img = similar(img, RGB{Float64})
  for y in 1:img_height, x in 1:img_width
    ix = vector_ix(img_width, y, x)
    result_img[y, x] = RGB{Float64}(solution[ix], solution[ix+1], solution[ix+2])
  end

  result_img
end

# Construct the system of equations.
function construct_system(
  img::Matrix{RGBA{Float64}},
  transform::AbstractAffineMap,
)
  img_height, img_width = size(img)

  num_subpixels = img_height * img_width * 3

  println("image: $(img_height) * $(img_width) * 3 = $(num_subpixels)")

  Is = Vector{Int64}()
  Js = Vector{Int64}()
  Vs = Vector{Float64}()

  right_side = Vector{Float64}()

  sizehint!(Is, num_subpixels * 3)
  sizehint!(Js, num_subpixels * 3)
  sizehint!(Vs, num_subpixels * 3)

  sizehint!(right_side, num_subpixels)

  row_ix = 1
  for y in 1:img_height, x in 1:img_width
    for pixel_equation in construct_pixel_equations(img, transform, y, x)
      for weighted_pixel in pixel_equation.left_side
        push!(Is, row_ix)
        push!(Js, weighted_pixel.vector_ix)
        push!(Vs, weighted_pixel.weight)
      end

      push!(right_side, pixel_equation.right_side)

      row_ix += 1
    end
  end

  println("Is: $(size(Is)), Js: $(size(Js)), Vs: $(size(Vs)), right_side: $(size(right_side))")

  matrix = sparse(Is, Js, Vs)
  println("matrix: $(size(matrix))")

  matrix, right_side
end

# Construct the equations corresponding to a single output pixel.
function construct_pixel_equations(
  img::Matrix{RGBA{Float64}},
  transform::AbstractAffineMap,
  y_int::Int64,
  x_int::Int64,
)
  y, x = Float64.((y_int, x_int))

  img_height, img_width = size(img)
  ix = vector_ix(img_width, y_int, x_int)

  # The polygon corresponding to the pixel under the transformation.
  corners_tr = Tuple.(transform.(SVector.([
    (y, x),
    (y, x + 1.0),
    (y + 1.0, x + 1.0),
    (y + 1.0, x)
  ])))

  # Determine the bounding box for the corresponding pixels following the
  # transformation.
  y_tr_min, y_tr_max = Int64.(floor.(extrema(coord[1] for coord in corners_tr)))
  x_tr_min, x_tr_max = Int64.(floor.(extrema(coord[2] for coord in corners_tr)))

  y_tr_min = max(y_tr_min, 1)
  x_tr_min = max(x_tr_min, 1)
  y_tr_max = min(y_tr_max, img_height)
  x_tr_max = min(x_tr_max, img_width)

  # Derivation of the equation for the pixel color:
  #
  # dst[y, x] =
  #   blend(
  #     src[y, x].alpha,
  #     sum(area * dst[y_tr, x_tr]) * multiply + add,
  #     src[y, x].color
  #    )
  #
  # let alpha = src[y, x].alpha
  # let beta = 1 - alpha
  #
  # dst[y, x] =
  #   beta * (sum(area * dst[y_tr, x_tr]) * multiply + add)
  #   + alpha * src[y, x].color
  #
  # dst[y, x] =
  #   beta * sum(area * dst[y_tr, x_tr]) * multiply
  #   + beta * add
  #   + alpha * src[y, x].color
  #
  # dst[y, x] - beta * sum(area * dst[y_tr, x_tr]) * multiply =
  #   beta * add + alpha * src[y, x].color
  #
  # dst[y, x] - sum(beta * area * dst[y_tr, x_tr] * multiply) =
  #   beta * add + alpha * src[y, x].color

  left_side_r = [WeightedPixel(vector_ix=ix + 0, weight=1.0)]
  left_side_g = [WeightedPixel(vector_ix=ix + 1, weight=1.0)]
  left_side_b = [WeightedPixel(vector_ix=ix + 2, weight=1.0)]

  img_pixel = img[y_int, x_int]
  beta = 1.0 - img_pixel.alpha

  right_side = beta * COLOR_ADD + img_pixel.alpha * convert(RGB{Float64}, img_pixel)

  if beta < SQRT_EPS
    # The remaining weights are all zero.
    return (
      PixelEquation(left_side=left_side_r, right_side=right_side.r),
      PixelEquation(left_side=left_side_g, right_side=right_side.g),
      PixelEquation(left_side=left_side_b, right_side=right_side.b),
    )
  end

  # The reciprocal of the area of the transformed pixel.
  recip_scale = 1.0 / area_polygon(corners_tr)

  # Add the rest of the left side terms by determining the contribution of each
  # transformed pixel.
  for y_tr_int in y_tr_min:y_tr_max, x_tr_int in x_tr_min:x_tr_max
    y_tr, x_tr = Float64.((y_tr_int, x_tr_int))

    ix_tr = vector_ix(img_width, y_tr_int, x_tr_int)

    # The intersection of the polygon corresponding to an actual pixel with the
    # transformed polygon.
    corners_pixel = [
      (y_tr, x_tr),
      (y_tr, x_tr + 1.0),
      (y_tr + 1.0, x_tr + 1.0),
      (y_tr + 1.0, x_tr)
    ]
    intersection = intersect_convex(corners_tr, corners_pixel)

    if isempty(intersection)
      continue
    end

    weight = (-beta * area_polygon(intersection) * recip_scale) * COLOR_MULTIPLY

    weighted_pixel_r = WeightedPixel(vector_ix=ix_tr + 0, weight=weight.r)
    weighted_pixel_g = WeightedPixel(vector_ix=ix_tr + 1, weight=weight.g)
    weighted_pixel_b = WeightedPixel(vector_ix=ix_tr + 2, weight=weight.b)

    push!(left_side_r, weighted_pixel_r)
    push!(left_side_g, weighted_pixel_g)
    push!(left_side_b, weighted_pixel_b)
  end

  (
    PixelEquation(left_side=left_side_r, right_side=right_side.r),
    PixelEquation(left_side=left_side_g, right_side=right_side.g),
    PixelEquation(left_side=left_side_b, right_side=right_side.b),
  )
end

comic()