Add better status tracking

1 files changed, 73 insertions, 25 deletions

diff --git a/gerbimg.py b/gerbimg.py
index 2d3cb72..51d2367 100755
--- a/gerbimg.py
+++ b/gerbimg.py
@@ -6,9 +6,11 @@ import tempfile
 import os.path as path
 import os
 import sys
+import time
 import shutil
 import math
 
+import tqdm
 import gerber
 from gerber.render import GerberCairoContext
 import numpy as np
@@ -21,6 +23,7 @@ def paste_image(
         subtract_gerber:list=[],
         extend_overlay_r_mil:float=12,
         extend_picture_r_mil:float=2,
+        status_print=lambda *args:None,
         debugdir:str=None):
     debugctr = 0
     def debugimg(img, name):
@@ -29,66 +32,99 @@ def paste_image(
             cv2.imwrite(path.join(debugdir, '{:02d}{}.png'.format(debugctr, name)), img)
         debugctr += 1
 
+    status_print('Parsing outline gerber')
     outline = gerber.loads(outline_gerber)
     (minx, maxx), (miny, maxy) = outline.bounds
     grbw, grbh = maxx - minx, maxy - miny
+    status_print('  * outline has offset {}, size {}'.format((minx, miny), (grbw, grbh)))
 
     imgh, imgw = source_img.shape
     scale = math.ceil(max(imgw/grbw, imgh/grbh)) # scale is in dpi
+    status_print('  * source image has size {}, going for scale {}dpi'.format((imgw, imgh), scale))
 
+    status_print('Parsing target gerber')
     target = gerber.loads(target_gerber)
     (tminx, tmaxx), (tminy, tmaxy) = target.bounds
+    status_print('  * target layer has offset {}, size {}'.format((tminx, tminy), (tmaxx-tminx, tmaxy-tminy)))
 
     with tempfile.TemporaryDirectory() as tmpdir:
         img_file = path.join(tmpdir, 'target.png')
 
+        status_print('Combining keepout composite')
         fg, bg = gerber.render.RenderSettings((1, 1, 1)), gerber.render.RenderSettings((0, 0, 0))
         ctx = GerberCairoContext(scale=scale)
+        status_print('  * target layer')
         ctx.render_layer(target, settings=fg, bgsettings=bg)
+        status_print('  * outline')
         ctx.render_layer(outline, settings=fg, bgsettings=bg)
-        for sub in subtract_gerber:
+        for i, sub in enumerate(subtract_gerber):
+            status_print('  * extra layer', i)
             layer = gerber.loads(sub)
             ctx.render_layer(layer, settings=fg, bgsettings=bg)
+        status_print('Rendering keepout composite')
         ctx.dump(img_file)
 
         original_img = cv2.imread(img_file, cv2.IMREAD_GRAYSCALE)
 
+    status_print('Expanding keepout composite')
     r = 1+2*max(1, int(extend_overlay_r_mil/1000 * scale))
     target_img = cv2.blur(original_img, (r, r))
     _, target_img = cv2.threshold(target_img, 255//(1+r), 255, cv2.THRESH_BINARY)
 
+    status_print('Thresholding source image')
     qr = 1+2*max(1, int(extend_picture_r_mil/1000 * scale))
     source_img = source_img[::-1]
     _, source_img = cv2.threshold(source_img, 127, 255, cv2.THRESH_BINARY)
     debugimg(source_img, 'thresh')
+
+    status_print('Padding source image')
     tgth, tgtw = target_img.shape
     padded_img = np.zeros(shape=(max(imgh, tgth), max(imgw, tgtw)), dtype=source_img.dtype)
-
     offx = int((minx-tminx if tminx < minx else 0)*scale)
     offy = int((miny-tminy if tminy < miny else 0)*scale)
     offx += int(grbw*scale - imgw) // 2
     offy += int(grbh*scale - imgh) // 2
     padded_img[offy:offy+imgh, offx:offx+imgw] = source_img
-
     debugimg(padded_img, 'padded')
+
+    status_print('Padding target image')
+    padded_target = np.zeros(shape=padded_img.shape, dtype=source_img.dtype)
+    offx = int(max(tminx, 0)*scale)
+    offy = int(max(tminy, 0)*scale)
+    padded_target[offy:offy+tgth, offx:offx+tgtw] = target_img
     debugimg(target_img, 'target')
-    out_img = (np.multiply((padded_img/255.0), (target_img/255.0) * -1 + 1) * 255).astype(np.uint8)
+    debugimg(padded_target, 'target_padded')
+
+    status_print('Masking source image')
+    out_img = (np.multiply((padded_img/255.0), (padded_target/255.0) * -1 + 1) * 255).astype(np.uint8)
 
     debugimg(out_img, 'multiplied')
-    debugimg(out_img + original_img, 'vis')
 
-    plot_contours(out_img, target, offx=(min(tminx, minx), min(tminy, miny)), scale=scale)
+    status_print('Calculating contour lines')
+    plot_contours(out_img,
+            target,
+            offx=(min(tminx, minx), min(tminy, miny)),
+            scale=scale,
+            status_print=lambda *args: status_print('   ', *args))
 
+    status_print('Generating output gerber')
     from gerber.render import rs274x_backend
     ctx = rs274x_backend.Rs274xContext(target.settings)
     target.render(ctx)
     return ctx.dump().getvalue()
 
 
-def plot_contours(img:np.ndarray, layer:gerber.rs274x.GerberFile, offx:tuple, scale:float, debug=lambda *args:None):
+def plot_contours(
+        img:np.ndarray,
+        layer:gerber.rs274x.GerberFile,
+        offx:tuple,
+        scale:float,
+        debug=lambda *args:None,
+        status_print=lambda *args:None):
     imgh, imgw = img.shape
 
     # Extract contours
+    status_print('Extracting contours')
     img_cont_out, contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_TC89_KCOS)
 
     aperture = list(layer.apertures)[0]
@@ -108,21 +144,26 @@ def plot_contours(img:np.ndarray, layer:gerber.rs274x.GerberFile, offx:tuple, sc
     done = []
     process_stack = [-1]
     next_process_stack = []
+    parents = [ (i, parent) for i, (_1, _2, _3, parent) in enumerate(hierarchy[0]) ]
     is_dark = True
-    while len(done) != len(contours):
-        for i, (_1, _2, _3, parent) in enumerate(hierarchy[0]):
-            if parent in process_stack:
-                contour = contours[i]
-                polarity = 'dark' if is_dark else 'clear'
-                debug('rendering {} with parent {} as {} with {} vertices'.format(i, parent, polarity, len(contour)))
-                debug('process_stack is', process_stack)
-                debug()
-                layer.primitives.append(Region(contour_lines(contour[:,0]), level_polarity=polarity, units=layer.settings.units))
-                next_process_stack.append(i)
-                done.append(i)
-        debug('skipping to next level')
-        process_stack, next_process_stack = next_process_stack, []
-        is_dark = not is_dark
+    status_print('Converting contours to gerber primitives')
+    with tqdm.tqdm(total=len(contours)) as progress:
+        while len(done) != len(contours):
+            for i, parent in parents[:]:
+                if parent in process_stack:
+                    contour = contours[i]
+                    polarity = 'dark' if is_dark else 'clear'
+                    debug('rendering {} with parent {} as {} with {} vertices'.format(i, parent, polarity, len(contour)))
+                    debug('process_stack is', process_stack)
+                    debug()
+                    layer.primitives.append(Region(contour_lines(contour[:,0]), level_polarity=polarity, units=layer.settings.units))
+                    next_process_stack.append(i)
+                    done.append(i)
+                    parents.remove((i, parent))
+                    progress.update(1)
+            debug('skipping to next level')
+            process_stack, next_process_stack = next_process_stack, []
+            is_dark = not is_dark
     debug('done', done)
 
 # Utility foo
@@ -167,13 +208,13 @@ def replace_file_in_zip(zip_path, filename, contents):
             zipout.writestr(filename, contents)
         shutil.move(tempname, zip_path)
 
-def paste_image_file(zip_or_dir, target, outline, source_img, subtract=[], debugdir=None):
+def paste_image_file(zip_or_dir, target, outline, source_img, subtract=[], status_print=lambda *args:None, debugdir=None):
     if path.isdir(zip_or_dir):
         tname, target = find_gerber_in_dir(zip_or_dir, target)
         _, outline = find_gerber_in_dir(zip_or_dir, outline)
         subtract = [ layer for _fn, layer in (find_gerber_in_dir(zip_or_dir, elem) for elem in subtract) ]
         
-        out = paste_image(target, outline, source_img, subtract, debugdir=debugdir)
+        out = paste_image(target, outline, source_img, subtract, debugdir=debugdir, status_print=status_print)
 
         # XXX
         with open('/tmp/out.GTO', 'w') as f:
@@ -183,7 +224,7 @@ def paste_image_file(zip_or_dir, target, outline, source_img, subtract=[], debug
         tname, target = find_gerber_in_zip(zip_or_dir, target)
         _, outline = find_gerber_in_zip(zip_or_dir, outline)
         
-        out = paste_image(target, outline, source_img, debugdir=debugdir)
+        out = paste_image(target, outline, source_img, subtract, debugdir=debugdir, status_print=status_print)
         replace_file_in_zip(zip_or_dir, tname, out)
     else:
         raise ValueError('{} does not look like either a folder or a zip file')
@@ -203,5 +244,12 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     source_img = cv2.imread(args.source, cv2.IMREAD_GRAYSCALE)
-    paste_image_file(args.zip_or_dir, args.target, args.outline, source_img, args.subtract, args.debug)
+    paste_image_file(
+            args.zip_or_dir,
+            args.target,
+            args.outline,
+            source_img,
+            args.subtract,
+            status_print=lambda *args: print(*args, flush=True),
+            debugdir=args.debug)