pre-crop

2 files changed, 14 insertions, 6 deletions

diff --git a/crop_frames.py b/crop_frames.py
index c22adb4..b507a1f 100644
--- a/crop_frames.py
+++ b/crop_frames.py
@@ -9,6 +9,7 @@ from PIL import Image
 parser = argparse.ArgumentParser()
 parser.add_argument('--step', default=256, type=int)
 parser.add_argument('--overlap', default=2, type=int)
+parser.add_argument('--ratio', default=2, type=int)
 parser.add_argument('--in_dir', help='Directory to process')
 parser.add_argument('--out_dir', default='inputs', help='Directory to output to')
 args = parser.parse_args()
@@ -21,7 +22,13 @@ def crop_dir():
   # print(files)
 
   img = Image.open(files[0])
-  width, height = img.size
+  width, old_height = img.size
+  if ratio == (width / height):
+    old_height = height
+  else:
+    height = round(width / ratio)
+    top_offset = round((old_height - height) / 2)
+
   print("{}x{} {}".format(width, height, step))
   for x in range(0, width, step):
     for y in range(0, height, step):
@@ -36,13 +43,13 @@ def crop_dir():
 
   dataset = []
   for raw_fn in files:
-    dataset.append((raw_fn, width, height, step, side, overlap,))
+    dataset.append((raw_fn, width, height, step, side, overlap, top_offset,))
 
   chunksize = 3 
   with Pool(processes=cpu_count()) as pool:
     pool.starmap(crop_image, dataset, chunksize)
 
-def crop_image(raw_fn, width, height, step, side, overlap):
+def crop_image(raw_fn, width, height, step, side, overlap, top_offset):
   img = Image.open(raw_fn)
   fn = os.path.basename(raw_fn)
   for x in range(0, width, step):
@@ -53,7 +60,7 @@ def crop_image(raw_fn, width, height, step, side, overlap):
       h = min(side, height - yy)
       crop_dir = "{}/crop_{}_{}_{}_{}".format(args.out_dir, x, y, w, h)
       # print("{}x{} {}x{}".format(x, y, w, h))
-      crop = img.crop((xx, yy, xx + w, yy + h))
+      crop = img.crop((xx, yy + top_offset, xx + w, yy + h + top_offset))
       crop.save("{}/{}".format(crop_dir, fn))
 
 if __name__ == "__main__":
diff --git a/test.py b/test.py
index eab51b2..13c92ef 100644
--- a/test.py
+++ b/test.py
@@ -16,6 +16,7 @@ from nets import FR_16L, FR_28L, FR_52L
 parser = argparse.ArgumentParser()
 parser.add_argument('--L', metavar='L', type=int, default=choice([16, 28, 52]), help='Network depth: One of 16, 28, 52')
 parser.add_argument('--T', metavar='T', default='L', help='Input type: L(Low-resolution) or G(Ground-truth)')
+parser.add_argument('--dataset', default=None, help='Name of dataset')
 parser.add_argument('--in_dir', metavar='in_dir', default=None, help='Directory to process')
 parser.add_argument('--out_dir', metavar='out_dir', default='/media/blue/uprez/results', help='Directory to output to')
 parser.add_argument('--network_dir', default='.', help='Path to networks')
@@ -43,10 +44,10 @@ if not(args.T == 'L' or args.T =='G'):
 
 def process_dir(dir):
     dir_partz = dir.split('/')
-    dataset = dir_partz[-2]
+    dataset = args.dataset = dir_partz[-2]
     part = dir_partz[-1]
     tag = '_'.join([dataset, str(args.L) + 'L', part])
-    out_path = os.path.join(args.out_dir, dataset, str(args.L) + 'L', part)
+    out_path = os.path.join(args.out_dir, dataset, part)
     render_path = os.path.join(args.out_dir, 'renders')
     os.makedirs(out_path, exist_ok=True)
     if args.mov_from_dirs: