Circles with Threads

README.md

Forks my own earlier gist: https://gist.github.com/jsundram/83685591e3e11a6ef2d2f7c4bc1aa6a6 to add Threading.

Run as follows: $ time julia -t 8 circles-threaded.jl

Output will look like:

Using 8 threads!
Getting Args...
No args supplied, using defaults: Dict{String,Any}("min_radius" => 2,"image" => "test.jpg")
  1.710215 seconds (2.22 M allocations: 111.929 MiB, 2.34% gc time)
Loading test.jpg ...
  1.593509 seconds (5.29 M allocations: 299.015 MiB, 5.62% gc time)
Resizing ...
  0.387044 seconds (1.48 M allocations: 93.052 MiB, 3.16% gc time)
Segmenting ...
  0.963858 seconds (2.64 M allocations: 243.068 MiB, 1.82% gc time)
Circling ...
	Thread 8 worked for 0.6871027946472168 seconds to produce 325 circles [473.0005503279413 c/s]
	Thread 7 worked for 3.0919532775878906 seconds to produce 610 circles [197.28629291445054 c/s]
	Thread 6 worked for 3.2021660804748535 seconds to produce 680 circles [212.35625601878897 c/s]
	Thread 4 worked for 8.612055778503418 seconds to produce 871 circles [101.13729200106972 c/s]
	Thread 5 worked for 8.758408784866333 seconds to produce 830 circles [94.76607228406125 c/s]
	Thread 2 worked for 25.866625547409058 seconds to produce 1520 circles [58.76298001121571 c/s]
	Thread 3 worked for 34.91539216041565 seconds to produce 1153 circles [33.02268508692798 c/s]
	Thread 1 worked for 79.90222191810608 seconds to produce 1853 circles [23.190844453602168 c/s]
make_circles produced 7843 circles in 80.38775897026062 seconds 97.56460561242304 c/s
 81.318932 seconds (5.19 M allocations: 66.456 GiB, 7.60% gc time)
Made 7843 circles
Drawing ...
  0.699790 seconds (189.30 k allocations: 12.377 MiB, 0.36% gc time)
julia -t 8 circles-threaded.jl  166.97s user 6.89s system 174% cpu 1:39.35 total

circles-threaded.jl

import ArgParse
import ImageSegmentation
import Images
import Luxor
import Statistics


int_round(f) = round(Int, f)


struct Circle{T}
    x::Int
    y::Int
    rgb::T
    r::Float64
end

# Is there a more normal way to do a keyword constructor for a type?
function Circle(;x::Int, y::Int, rgb::Tuple{Int64, Int64, Int64}, r::Float64)
    return Circle(x, y, rgb, r)
end


function parse_commandline()
    s = ArgParse.ArgParseSettings()
    s.exc_handler = ArgParse.debug_handler

    @ArgParse.add_arg_table s begin
        "--image", "-i"
            help = "Path to input image"
            arg_type = String
            required = true
        "--min_radius", "-r"
            help = "Minimum radius of generated circles"
            arg_type = Float64
            default = 2.0
    end

    return ArgParse.parse_args(s)
end


function scale(img, target_px)
    h, w = size(img)
    # Need sqrt because ratio gets applied in both x and y dimension.
    return Images.imresize(img, ratio=(target_px / (h*w))^.5)
end


function segment(img)
    # k controls segment size. 5 - 500 is a good range
    # min_size gets rid of small segments 
    return ImageSegmentation.felzenszwalb(img, 60, 100)
end


function make_mask(label_ix, imgh, imgw)
    # Create a mask array that's just large enough to contain
    # the segment we are working on. This will allow the distance transform
    # to operate the smallest array possible, for speed.
    pad = Dict(true => 0, false => 1)
    (top, left), (bottom, right) = map(ix -> ix.I, extrema(label_ix))

    # Allocate a pixel-wide boundary strip if we aren't at the image border.
    top_pad, bottom_pad = pad[top == 1], pad[bottom == imgh]
    left_pad, right_pad = pad[left == 1], pad[right == imgw]

    mask = ones(Bool,
        bottom - top + 1 + top_pad + bottom_pad,
        right - left + 1 + left_pad + right_pad
    )

    new_ix = map(x -> x - CartesianIndex(top - top_pad - 1, left - left_pad - 1), label_ix)
    @inbounds mask[new_ix] .= 0

    return (left - left_pad, top - top_pad, mask)
end


function circle_indices(center, radius, height, width)
    # Code inspired by:
    # https://github.com/JuliaImages/ImageDraw.jl/blob/master/src/ellipse2d.jl

    indices = CartesianIndex{2}[]
    r = CartesianIndex(int_round(radius), int_round(radius))
    y, x = center.I
    # Examine (2*radius + 1)^2 spots, expecting array length ~ pi * radius^2.
    for ix in center - r: center + r
        row, col = ix.I
        if 1 <= row <= height && 1 <= col <= width
            if (row - y)^2 + (col - x)^2 <= radius^2
                push!(indices, ix)
            end
        end
    end

    return indices
end


function make_circles(img, seg, min_radius)
    float2int(f) = round(Int, 256*f)
    to_clr(l) = float2int.( (l.r, l.g, l.b) )

    # Arrays for returning values from threads in. keyed on threadid to know
    # we won't have race conditions or contentions for writing.
    D = Dict([(ix, []) for ix in 1:Threads.nthreads()])
    # T (times) and C (circles) are just for keeping track of thread work stats.
    T = Dict([(ix, 0.0) for ix in 1:Threads.nthreads()])
    C = Dict([(ix, 0) for ix in 1:Threads.nthreads()])
    
    # Set up an array of jobs in such a way that the labor is a bit more fairly
    # distributed (depends on reverse engineering how :static works)
    ml = [findall(x -> x == l, seg.image_indexmap) for l in seg.segment_labels]
    ml = sort(ml, by=length, rev=true)
    n = Threads.nthreads()
    jobs = []
    for i in 1:n
        append!(jobs, ml[i:n:length(ml)])
    end

    h, w = size(img)
    Threads.@threads for m in jobs
        start_time = time()
        left, top, mask = make_mask(m, h, w)

        my_circles = []
        while true
            f = Images.feature_transform(mask)
            edt = Images.distance_transform(f)
            ix = argmax(edt)
            r = edt[ix]

            if r < min_radius || r == Inf
                break
            end

            circle_ix = circle_indices(ix, r, size(edt)...)
            @inbounds mask[circle_ix] .= 1

            offset_ix = map(x -> x + CartesianIndex(top - 1, left - 1), circle_ix)
            clr = to_clr(Statistics.mean(img[offset_ix]))

            y, x = ix.I

            push!(my_circles, Circle(
                x=left + x,
                y=top + y,
                rgb=clr,
                r=r
            ))
        end
        run_time = time() - start_time

        # Save our circles!
        id = Threads.threadid()
        append!(D[id], my_circles)

        # Update stats
        T[id] += run_time
        C[id] += length(my_circles)
        # println("Thread $(id) produced $(length(my_circles)) circles in $(run_time) seconds [$(length(my_circles) / run_time) c/s]")
    end

    # Add a big circle with the avg color for the background.
    circles = [Circle(
        x=int_round(w/2),
        y=int_round(h/2),
        rgb= to_clr(Statistics.mean(img)),
        r=(w*w + h*h)^.5 + 2
    )]

    # Collate computed circles
    for (threadid, cl) in D
        append!(circles, cl)
    end

    # Display per-thread stats (an indication of how well spread-out the work was)
    for (threadid, f) in sort(collect(T), by=p->p[2])
        c = C[threadid]
        println("\tThread $(threadid) worked for $(f) seconds to produce $(c) circles [$(c / f) c/s]")
    end

    return circles
end


function draw(img, circles, filename)
    h, w = size(img)
    Luxor.Drawing(w, h, filename)
    # Paint from biggest to smallest
    for c in sort(circles, by=c -> c.r, rev=true)
        r, g, b = map(x -> x / 255, c.rgb)
        Luxor.setcolor(r, g, b)
        Luxor.circle(c.x, c.y, c.r - .5, :fill)
    end
    Luxor.finish()
end


function get_args()
    args = nothing
    try
        args = parse_commandline()
        println("args", args)
    catch err
        args = Dict("image" => "test.jpg", "min_radius" => 2)
        println("No args supplied, using defaults: ", args)
    end
    return args
end 


function main()
    println("Using $(Threads.nthreads()) threads!")

    println("Getting Args...")
    @time args = get_args()

    path = args["image"]
    name, _ = splitext(basename(path))

    println("Loading $(path) ...")
    @time img = Images.load(path)

    println("Resizing ...")
    @time img = scale(img, 1e6)  # Downsample to ~1M pixels
    Images.save("./$(name)-orig.png", img)

    println("Segmenting ... ")
    @time seg = segment(img)
    Images.save("./$(name)-seg.png", map(i-> seg.segment_means[i], seg.image_indexmap))

    println("Circling ...")
    @time circles = make_circles(img, seg, args["min_radius"])
    println("Made $(length(circles)) circles")

    println("Drawing ...")
    @time draw(img, circles, "./$(name)-circled.png")
end


if !isinteractive()
    main()
end