change to cli_proc

1 files changed, 33 insertions, 5 deletions

diff --git a/megapixels/commands/demo/face_search.py b/megapixels/commands/demo/face_search.py
index f551cafd..4c7036f4 100644
--- a/megapixels/commands/demo/face_search.py
+++ b/megapixels/commands/demo/face_search.py
@@ -53,7 +53,7 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_results, opt_gpu):
   dataset.load_metadata(types.Metadata.FILE_RECORD)
   dataset.load_metadata(types.Metadata.FACE_VECTOR)
   dataset.load_metadata(types.Metadata.FACE_ROI)
-  # dataset.load_metadata(types.Metadata.IDENTITY)
+  dataset.load_metadata(types.Metadata.IDENTITY)
 
   # init face detection
   detector = face_detector.DetectorCVDNN()
@@ -82,22 +82,50 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_results, opt_gpu):
   image_records = dataset.find_matches(vec_query, n_results=opt_results)
 
   # summary
-  im_query = draw_utils.draw_bbox(im_query, bbox_norm, stroke_weight=8)
+  im_query = draw_utils.draw_bbox(im_query, bbox_norm, stroke_weight=4)
   ims_match = [im_query]
+
+  opt_size = (256,256)
+
   for image_record in image_records:
     image_record.summarize()
     log.info(f'{image_record.filepath}')
     im_match = cv.imread(image_record.filepath)
+    dim_match = im_match.shape[:2][::-1]
+    bbox_match = image_record.bbox
+    score = image_record.score
+    if score < .5:
+      clr = (0,255,0)
+    elif score < .6:
+      clr = (0,255,125)
+    elif score < .65:
+      clr = (0,125,125)
+    elif score < .7:
+      clr = (0,125,255)
+    else:
+      clr = (0,0,255)
+      
+    im_match = draw_utils.draw_bbox(im_match, bbox_match, stroke_weight=4, color=clr )
+    bbox_match_dim = bbox_match.to_dim(dim_match)
 
-    im_match_pil = Image.open(image_record.filepath).convert('RGB')
-    # bbox = 
-    ims_match.append(im_match)
+    im_pil = im_utils.ensure_pil(im_match)
+    center = (bbox_match_dim.cx, bbox_match_dim.cy)
+    im_pil = ImageOps.fit(im_pil, opt_size, centering=center)
+    im_np = im_utils.ensure_np(im_pil)
+    if image_record.identity is not None:
+      log.debug(f'identity: {image_record.identity.name_display}')
+    else:
+      log.debug('no identity info')
+    log.debug(f'score: {image_record.score}')
 
+    ims_match.append(im_np)
   # make montages of most similar faces
   montages = imutils.build_montages(ims_match, (256, 256), (3,2))
 
   # display 
   for i, montage in enumerate(montages):
     cv.imshow(f'{opt_dataset.name.upper()}: page {i}', montage)
+    fp_out = join(Path(opt_fp_in).parent, f'{Path(opt_fp_in).stem}_{i}.png')
+    cv.imwrite(fp_out, montage)
   
   display_utils.handle_keyboard()