augmentor: pass alphas as heatmaps

Signed-off-by: Niv Sardi <xaiki@evilgiggle.com>

python/augment.py

@@ -24,6 +24,8 @@ BATCH_SIZE = 16
 mkdir.make_dirs([defaults.AUGMENTED_IMAGES_PATH, defaults.AUGMENTED_LABELS_PATH])
 
 logo_images = []
+logo_alphas = []
+
 background_images = [d for d in os.scandir(defaults.IMAGES_PATH)]
 
 stats = {
@@ -57,14 +59,18 @@ for d in os.scandir(defaults.LOGOS_DATA_PATH):
 
         assert(w > 10)
         assert(h > 10)
+        (b, g, r, _) = cv2.split(img)
+        alpha = img[:, :, 3]/255
+        logo_images.append(cv2.merge([b, g, r]))
+        # XXX(xaiki): we pass alpha as a float32 heatmap, because imgaug is pretty strict about what data it will process
+        logo_alphas.append(np.dstack((alpha, alpha, alpha)).astype('float32'))
 
-        logo_images.append(img)
     except Exception as e:
         stats['failed'] += 1
         print(f'error loading: {d.path}: {e}')
 
 print(stats)
-batches = [UnnormalizedBatch(images=logo_images[i:i+BATCH_SIZE])
+batches = [UnnormalizedBatch(images=logo_images[i:i+BATCH_SIZE],heatmaps=logo_alphas[i:i+BATCH_SIZE])
            for i in range(math.floor(len(logo_images)/BATCH_SIZE))]
 
 # We use a single, very fast augmenter here to show that batches
@@ -91,12 +97,17 @@ with pipeline.pool(processes=-1, seed=1) as pool:
 
             anotations = []
             for k in range(math.floor(len(batch_aug.images_aug)/3)):
-                logo = batch_aug.images_aug[(j+k)%len(batch_aug.images_aug)]
+                logo_idx = (j+k*4)%len(batch_aug.images_aug)
+                logo = batch_aug.images_aug[logo_idx]
+
+                # XXX(xaiki): we get alpha from heatmap, but will only use one channel
+                # we could make mix_alpha into mix_mask and pass all 3 chanels
+                alpha = cv2.split(batch_aug.heatmaps_aug[logo_idx])
                 try:
-                    img, bb, (w, h) = imtool.mix(img, logo, random.random(), random.random())
+                    img, bb, (w, h) = imtool.mix_alpha(img, logo, alpha[0], random.random(), random.random())
                     anotations.append(f'0 {bb.x/w} {bb.y/h} {bb.w/w} {bb.h/h}')
-                except AssertionError:
-                    print(f'couldnt process {i}, {j}')
+                except AssertionError as e:
+                    print(f'couldnt process {i}, {j}: {e}')
 
             try:
                 cv2.imwrite(f'{defaults.AUGMENTED_IMAGES_PATH}/{basename}.png', img)

python/imtool.py

@@ -89,12 +89,32 @@ def remove_white(img):
 
     return rect
 
+
 def mix(a, b, fx, fy):
+    alpha = b[:, :, 3]/255
+
+    return _mix_alpha(a, b, alpha, fx, fy)
+
+def mix_alpha(a, b, ba, fx, fy):
     (ah, aw, ac) = a.shape
     (bh, bw, bc) = b.shape
 
-    assert(aw > bw)
-    assert(ah > bh)
+    if (aw < bw or ah < bh):
+        f = 0.2*aw/bw
+        print(f'resizing, factor {f} to fit in {aw}x{ah}\n -- {bw}x{bh} => {floor_point(bw*f, bh*f)}')
+        r = cv2.resize(b, floor_point(bw*f, bh*f), interpolation = cv2.INTER_LINEAR)
+        rba = cv2.resize(ba, floor_point(bw*f, bh*f), interpolation = cv2.INTER_LINEAR)
+
+        return mix_alpha(a, r, rba, fx, fy)
+
+    assert bw > 10, f'b({bw}) too small'
+    assert bh > 10, f'b({bh}) too small'
+
+    return _mix_alpha(a, b, ba, fx, fy)
+
+def _mix_alpha(a, b, ba, fx, fy):
+    (ah, aw, ac) = a.shape
+    (bh, bw, bc) = b.shape
 
     x = math.floor(fx*(aw - bw))
     y = math.floor(fy*(ah - bh))
@@ -102,8 +122,7 @@ def mix(a, b, fx, fy):
     # handle transparency
     mat = a[y:y+bh,x:x+bw]
     cols = b[:, :, :3]
-    alpha = b[:, :, 3]/255
-    mask = np.dstack((alpha, alpha, alpha))
+    mask = np.dstack((ba, ba, ba))
 
     a[y:y+bh,x:x+bw] = mat * (1 - mask) + cols * mask