python: fixups

Signed-off-by: Niv Sardi <[email protected]>

python/augment.py

@@ -140,9 +140,12 @@ with pipeline.pool(processes=-1, seed=1) as pool:
                 label = logo_labels[orig.tobytes()]
                 logo = batch_aug.images_aug[logo_idx]
 
+                assert(logo.shape == orig.shape)
+
                 # 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:
                     bb = imtool.mix_alpha(img, logo, alpha[0],
                                           random.random(), random.random())

python/entity.py

@@ -4,8 +4,12 @@ from typing import NamedTuple
 
 from common import defaults
 
-def read_entities(fn):
-    with open(defaults.MAIN_DATA_PATH, newline='') as csvfile:
+def update(d, i):
+    d.update(i)
+    return d
+
+def read_entities(fn = defaults.MAIN_CSV_PATH):
+    with open(fn, newline='') as csvfile:
         reader = csv.DictReader(csvfile)
         bcos = { d['bco']:update(d, {'id': i}) for i, d in enumerate(reader)}
     return bcos

python/imtool.py

@@ -21,7 +21,6 @@ class BoundingBox(NamedTuple):
     @classmethod
     def from_centroid(cls, c, shape):
         (ih, iw, ic) = shape
-        print(cls, c, shape)
         self = cls(x=math.floor(w*(c.x - c.w/2))
                    , y=math.floor(h*(c.y - c.h/2))
                    , w=math.ceil(w*c.w)
@@ -158,7 +157,7 @@ def mix_alpha(a, b, ba, fx, fy):
     if (aw*p < bw or ah*p < bh):
         f = min(p*aw/bw, p*ah/bh)
         nw, nh = floor_point(bw*f, bh*f)
-        print(f'resizing, factor {f} to fit in {aw}x{ah}\n -- {bw}x{bh} => {nw}x{nh}')
+        print(f'resizing to fit in {aw}x{ah}\t {bw}x{bh}\t=> {nw}x{nh}\tfactor {f}')
         r = cv2.resize(b, (nw, nh), interpolation = cv2.INTER_LINEAR)
         rba = cv2.resize(ba, (nw, nh), interpolation = cv2.INTER_LINEAR)
 
@@ -232,7 +231,7 @@ def crop(id, fn, logos: List[Centroid], out = './data/squares'):
                     li.append(p)
 
             nim = cut_img(im, start, end)
-            rnim = cut_img(im, start, end)
+            rnim = cut_img(rim, start, end)
             img_name =f"{img_out}/{basename}-x{x}y{y}.jpg"
             txt_name =f"{txt_out}/{basename}-x{x}y{y}.txt"