Options for remove_highlights script

bob/ip/qualitymeasure/__init__.py

@@ -6,6 +6,7 @@ from .msu_iqa_features import compute_msu_iqa_features
 
 from ._library import remove_highlights
 from ._library_orig import remove_highlights_orig
+from . import tan_specular_highlights as tsh
 
 
 def get_config():

bob/ip/qualitymeasure/script/remove_highlights.py

@@ -9,12 +9,15 @@ Created on 28 Jun 2017
 
 import sys
 import argparse
+import time
 
 import bob.io.base
 import bob.io.image
 import numpy as np
 
 from bob.ip.qualitymeasure import remove_highlights_orig
+from bob.ip.qualitymeasure import remove_highlights
+from bob.ip.qualitymeasure import tsh
 
 def main(command_line_parameters=None):
     """Remove the specular component of the input image and write result to
@@ -38,6 +41,27 @@ def main(command_line_parameters=None):
         default=None,
         help='filename of specular free image.')
 
+    argParser.add_argument(
+        '-r',
+        '--residue',
+        dest='resImg',
+        default=None,
+        help='filename of specular residue image.')
+
+    argParser.add_argument(
+        '-e',
+        '--epsilon',
+        dest='epsilon',
+        default=0.5,
+        help='value of epsilon parameter.')
+
+    argParser.add_argument(
+        '-a',
+        '--algorithm',
+        dest='algorithm',
+        default=0,
+        help='version of the algoritm used.')
+
     args = argParser.parse_args(command_line_parameters)
 
     if not args.inpImg:
@@ -52,11 +76,24 @@ def main(command_line_parameters=None):
     # 2. compute
     print("Extracting diffuse component...")
 
-    sfi, diff, residue = remove_highlights_orig(img.astype(np.float32), 0.5)
+    if int(args.algorithm) == 0:
+        print('v0')
+        sfi, diff, residue = tsh.remove_highlights(img.astype(np.float32), float(args.epsilon))
+    elif int(args.algorithm) == 1:
+        print('v1')
+        sfi, diff, residue = remove_highlights_orig(img.astype(np.float32), float(args.epsilon))
+    else:
+        print('v2')
+        sfi, diff, residue = remove_highlights(img.astype(np.float32), float(args.epsilon))
 
     # 1. save output image
     print("Saving output file: %s" % args.outImg)
     bob.io.base.save(diff.astype('uint8'), args.outImg)
+
+    if args.resImg:
+        print("Saving residue file: %s" % args.resImg)
+        bob.io.base.save(residue.astype('uint8'), args.resImg)
+
     print('Done')