Diff of /scripts/subsample.py [000000] .. [5a7589]

Switch to side-by-side view

--- a
+++ b/scripts/subsample.py
@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+import os
+import sys
+import json
+import numpy as np
+from PIL import Image
+from pycocotools.mask import encode, decode
+
+def get_outfile(infile, outdir):
+    outfile = os.path.basename(infile)
+    outsubdir = os.path.basename(os.path.dirname(infile))
+    outfile = os.path.join(outdir, outsubdir, outfile)
+    return outfile
+
+
+def filegen(indir, ext='png'):
+    for dd in os.scandir(indir):
+        if not os.path.isdir(dd.path):
+            continue
+        for ff in os.scandir(dd.path):
+            if not ff.name.endswith(ext):
+                continue
+            yield ff.path
+
+
+def subsample_coco_mask(mask, factor):
+    mask = Image.fromarray(mask)
+    sz = np.asarray(mask.size)
+    out_sz = (sz / factor).astype(int)
+    mask.thumbnail(size, Image.ANTIALIAS)
+    mask = np.asarray(mask, order='F')
+    return mask
+
+
+def subsample_verts(verts, factor):
+    verts = (np.asarray(verts)//factor).tolist()
+
+    vvprev = [None, None]
+
+    newverts = []
+    for vv in verts:
+        if np.all(vv == vvprev):
+           continue
+        else:
+            newverts.append(vv)
+        vvprev = vv.copy()
+    return newverts
+
+
+def subsample_roi(fn, fnout, factor):
+    with open(fn) as fh:
+        rois = json.load(fh)
+
+    if len(rois) == 0:
+        print("NO ROIS IN\t%s" % fn)
+
+    for roi in rois:
+        mask = decode(roi)
+        mask = subsample_coco_mask(mask, factor)
+        cocomask = encode(mask)
+        cocomask['counts'] = cocomask['counts'].decode()
+        roi.update(cocomask)
+
+        verts = roi["vertices"]
+        roi.update({"vertices": subsample_verts(verts, factor)})
+        if 'zoom' in roi:
+            roi.pop('zoom')
+
+    with open(fnout, 'w+') as fh:
+        json.dump(rois, fh)
+
+####################################################################
+if __name__ == '__main__':
+    import sys
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        'indir', type=str, help = 'input directory')
+
+    parser.add_argument(
+        '--outdir',
+        default='',
+        type=str,
+        help = 'output directory')
+
+    parser.add_argument(
+      '--original-side',
+      type=int,
+      default=1024,
+      help='Original side (in pixels) of the image patches')
+
+
+    parser.add_argument('--img', dest='img', action='store_true')
+    parser.add_argument('--no-img', dest='img', action='store_false')
+    parser.set_defaults(img=True)
+
+    group = parser.add_mutually_exclusive_group(required=True)
+
+    group.add_argument(
+      '--target-side',
+      type=int,
+      default=512,
+      help='Requested side (in pixels) of the image patches')
+
+    group.add_argument(
+      '--factor',
+      type=int,
+      default=None,
+      help='Requested subsampling rate')
+    
+    prms = parser.parse_args()
+
+    if prms.factor is None:
+        factor = prms.original_side // prms.target_side 
+        side = prms.target_side
+    else:
+        factor = prms.factor
+        side = prms.original_side // prms.factor
+    size = side, side
+
+    if len(prms.outdir)==0:
+        basedir, set_ = os.path.split(prms.indir.rstrip('/'))
+        basedir = os.path.dirname(basedir)
+        #basedir = os.path.dirname(basedir)
+        outdir = "{}/data_{}_subsample_{:d}x/fullsplit/all".format(basedir, side, factor)
+    else:
+        outdir = prms.outdir
+
+    print("SAVING TO:", outdir, sep='\t')
+    os.makedirs(outdir, exist_ok = True)
+    
+    if prms.img:
+        print("SUBSAMPLING IMAGES")
+        for infile in filegen(prms.indir):
+            outfile = get_outfile(infile, outdir=outdir)
+            os.makedirs(os.path.dirname(outfile), exist_ok=True)
+            if infile != outfile:
+                try:
+                    im = Image.open(infile)
+                    im.thumbnail(size, Image.ANTIALIAS)
+                    im.save(outfile, "png")
+                except IOError as ee:
+                    print( "cannot create thumbnail for '%s'" % infile)
+                    print(ee)
+
+    print("SUBSAMPLING MASKS")
+    for infile in filegen(prms.indir, ext = '.json'):
+        outfile = get_outfile(infile, outdir=outdir)
+        if infile != outfile:
+            try:
+                subsample_roi(infile, outfile, factor)
+            except IOError as ee:
+                print( "cannot create thumbnail for '%s'" % infile)
+                print(ee)
+            except Exception as ee:
+                print( "cannot create thumbnail for '%s'" % infile)
+                raise ee