Fix batches > 1

.github/workflows/CompatHelper.yml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: [1.2.0]
+        julia-version: [1.6]
         julia-arch: [x86]
         os: [ubuntu-latest]
     steps:

.github/workflows/RunTests.yml

@@ -8,15 +8,15 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        julia-version: ['1', '^1.6.0-0', 'nightly']
+        julia-version: ['1.6', 'nightly']
         julia-arch: [x64]
         os: [ubuntu-latest, macOS-latest, windows-latest]
         exclude:
           - os: macOS-latest
             julia-arch: x86
 
     steps:
-      - uses: actions/checkout@v1.0.0
+      - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@latest
         with:
           version: ${{ matrix.julia-version }}

src/yolo/yolo.jl

@@ -625,18 +625,20 @@ function (yolo::yolo)(img::DenseArray; detectThresh=nothing, overlapThresh=yolo.
         end
     end
 
-
     classes = unique(batchout[end-1, :])
-    output = Array{Array{Float32, 2},1}(undef, 0)
-    for c in classes
-        detection = sortslices(batchout[:, batchout[end-1, :] .== c], dims = 2, by = x -> x[5], rev = true)
-        for l in 1:size(detection, 2)
-            iou = bboxiou(view(detection, 1:4, l), detection[1:4, l+1:end])
-            ds = findall(v -> v >= overlapThresh, iou)
-            detection = detection[:, setdiff(1:size(detection, 2), ds .+ l)]
-            l >= size(detection,2) && break
+    output = Array{Float32, 2}[]
+    @views for b in 1:yolo.cfg[:batchsize]
+        page = batchout[:, batchout[end,:] .== b]
+        for c in classes
+            detection = sortslices(page[:, page[end-1, :] .== c], dims = 2, by = x -> x[5], rev = true)
+            for det_idx in 1:size(detection, 2)
+                iou = bboxiou(detection[1:4, det_idx], detection[1:4, det_idx+1:end])
+                same_objects = findall(>=(overlapThresh), iou)
+                detection = detection[:, setdiff(1:size(detection, 2), same_objects .+ det_idx)]
+                det_idx >= size(detection,2) && break
+            end
+            push!(output, detection)
         end
-        push!(output, detection)
     end
     return hcat(output...)
 end

test/maintests.jl

@@ -22,6 +22,41 @@ pretrained_list = [
                     # YOLO.v3_spp_608_COCO
                     ]
 
+@testset "batch sizes" begin
+    IMG = load(joinpath(@__DIR__, "images", "dog-cycle-car.png"))
+    @testset "Batch size $batch_size" for batch_size in [1, 3]
+        yolomod = YOLO.v3_tiny_416_COCO(batch = batch_size, silent=true)
+        batch = emptybatch(yolomod)
+        @test size(batch) == (416, 416, 3, batch_size)
+        for b in 1:batch_size
+            batch[:, :, :, b], padding = prepareImage(IMG, yolomod)
+        end
+        res = yolomod(batch, detectThresh  = dThresh, overlapThresh = oThresh);
+        @test size(res) == (89, 4 * batch_size)
+        for b in 2:batch_size
+            @test res[1:end-1, res[end,:] .== 1] == res[1:end-1, res[end,:] .== b]
+        end
+    end
+end
+
+@testset "Custom cfg's" begin
+    @testset "Valid non-square dimensions (512x384)" begin
+        IMG = load(joinpath(@__DIR__,"images","dog-cycle-car.png"))
+        yolomod = YOLO.v3_COCO(silent=true, cfgchanges=[(:net, 1, :width, 512), (:net, 1, :height, 384)])
+        batch = emptybatch(yolomod)
+        batch[:,:,:,1], padding = prepareImage(IMG, yolomod)
+        res = yolomod(batch, detectThresh=dThresh, overlapThresh=oThresh) #run once
+        @test size(res,2) > 0
+    end
+    @testset "Invalid non-square dimensions" begin
+        IMG = load(joinpath(@__DIR__,"images","dog-cycle-car.png"))
+        # invalid height
+        @test_throws AssertionError YOLO.v3_COCO(silent=false, w=512, h=383)
+        # invalid width
+        @test_throws AssertionError YOLO.v3_COCO(silent=false, w=511, h=384)
+    end
+end
+
 header = ["Model" "loaded?" "load time (s)" "ran?" "run time (s)" "objects detected"]
 table = Array{Any}(undef, length(pretrained_list), 6)
 for (k, pretrained) in pairs(pretrained_list)
@@ -65,22 +100,3 @@ for (k, pretrained) in pairs(pretrained_list)
 end
 pretty_table(table, header)
 @info "Times approximate. For more accurate benchmarking run ObjectDetector.benchmark()"
-
-
-@testset "Custom cfg's" begin
-    @testset "Valid non-square dimensions (512x384)" begin
-        IMG = load(joinpath(@__DIR__,"images","dog-cycle-car.png"))
-        yolomod = YOLO.v3_COCO(silent=true, cfgchanges=[(:net, 1, :width, 512), (:net, 1, :height, 384)])
-        batch = emptybatch(yolomod)
-        batch[:,:,:,1], padding = prepareImage(IMG, yolomod)
-        res = yolomod(batch, detectThresh=dThresh, overlapThresh=oThresh) #run once
-        @test size(res,2) > 0
-    end
-    @testset "Invalid non-square dimensions" begin
-        IMG = load(joinpath(@__DIR__,"images","dog-cycle-car.png"))
-        # invalid height
-        @test_throws AssertionError YOLO.v3_COCO(silent=false, w=512, h=383)
-        # invalid width
-        @test_throws AssertionError YOLO.v3_COCO(silent=false, w=511, h=384)
-    end
-end