externalize drawing, cleanup

1 files changed, 20 insertions, 22 deletions

diff --git a/megapixels/commands/cv/face_roi.py b/megapixels/commands/cv/face_roi.py
index c3c2ac05..6d42924e 100644
--- a/megapixels/commands/cv/face_roi.py
+++ b/megapixels/commands/cv/face_roi.py
@@ -29,7 +29,7 @@ color_filters = {'color': 1, 'gray': 2, 'all': 3}
   show_default=True,
   help=click_utils.show_help(types.Dataset))
 @click.option('--size', 'opt_size', 
-  type=(int, int), default=(300, 300),
+  type=(int, int), default=(480, 480),
   help='Output image size')
 @click.option('-d', '--detector', 'opt_detector_type',
   type=cfg.FaceDetectNetVar,
@@ -50,7 +50,7 @@ color_filters = {'color': 1, 'gray': 2, 'all': 3}
 @click.option('--color', 'opt_color_filter', 
   type=click.Choice(color_filters.keys()), default='all',
   help='Filter to keep color or grayscale images (color = keep color')
-@click.option('--largest/--all-faces', 'opt_largest', is_flag=True, default=True,
+@click.option('--keep', 'opt_largest', type=click.Choice(['largest', 'all']), default='all',
   help='Only keep largest face')
 @click.option('--zone', 'opt_zone', default=(0.0, 0.0), type=(float, float), 
   help='Face center must be located within zone region (0.5 = half width/height)')
@@ -72,7 +72,7 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset,
   import cv2 as cv
   import pandas as pd
 
-  from app.utils import logger_utils, file_utils, im_utils
+  from app.utils import logger_utils, file_utils, im_utils, display_utils, draw_utils
   from app.processors import face_detector
   from app.models.data_store import DataStore
 
@@ -113,13 +113,15 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset,
 
   # filter out grayscale
   color_filter = color_filters[opt_color_filter]
+  # set largest flag, to keep all or only largest
+  opt_largest = opt_largest == 'largest'
 
   data = []
 
   for df_record in tqdm(df_records.itertuples(), total=len(df_records)):
     fp_im = data_store.face(str(df_record.subdir), str(df_record.fn), str(df_record.ext))
     im = cv.imread(fp_im)
-
+    im_resized = im_utils.resize(im, width=opt_size[0], height=opt_size[1])
     # filter out color or grayscale iamges
     if color_filter != color_filters['all']:
       try:
@@ -130,9 +132,10 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset,
       except Exception as e:
         log.error('Could not check grayscale: {}'.format(fp_im))
         continue
-        
+    
     try:
-      bboxes = detector.detect(im, size=opt_size, pyramids=opt_pyramids, largest=opt_largest, zone=opt_zone)
+      bboxes = detector.detect(im_resized, pyramids=opt_pyramids, largest=opt_largest, 
+        zone=opt_zone, conf_thresh=opt_conf_thresh)
     except Exception as e:
       log.error('could not detect: {}'.format(fp_im))
       log.error('{}'.format(e))
@@ -150,27 +153,22 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset,
     if len(bboxes) == 0:
       log.warn(f'no faces in: {fp_im}')
     
-    # debug display
+    # if display optined
     if opt_display and len(bboxes):
-      im_md = im_utils.resize(im, width=min(1200, opt_size[0]))
+      # draw each box
       for bbox in bboxes:
-        bbox_dim = bbox.to_dim(im_md.shape[:2][::-1])
-        cv.rectangle(im_md, bbox_dim.pt_tl, bbox_dim.pt_br, (0,255,0), 3)
-      cv.imshow('', im_md)
-      while True:
-        k = cv.waitKey(1) & 0xFF
-        if k == 27 or k == ord('q'):  # ESC
-          cv.destroyAllWindows()
-          sys.exit()
-        elif k != 255:
-          # any key to continue
-          break
+        bbox_dim = bbox.to_dim(im_resized.shape[:2][::-1])
+        draw_utils.draw_bbox(im_resized, bbox_dim)
 
-  # save date
+      # display and wait
+      cv.imshow('', im_resized)
+      display_utils.handle_keyboard()
+
+  # create DataFrame and save to CSV
   file_utils.mkdirs(fp_out)
   df = pd.DataFrame.from_dict(data)
   df.index.name = 'index'
   df.to_csv(fp_out)
+
   # save script
-  cmd_line = ' '.join(sys.argv)
-  file_utils.write_text(cmd_line, '{}.sh'.format(fp_out))
\ No newline at end of file
+  file_utils.write_text(' '.join(sys.argv), '{}.sh'.format(fp_out))
\ No newline at end of file