Misc bug fixes to example scripts

abhinaukumar · Oct 31, 2023 · 070a450 · 070a450
1 parent 7f696ca
commit 070a450
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 100 deletions.
diff --git a/examples/extract_features.py b/examples/extract_features.py
@@ -1,5 +1,7 @@
 import argparse
 
+from feature_extractors.ssim_fex import SsimFeatureExtractor  # Import feature extractor(s) to make visible to get_fex
+from qualitylib.feature_extractor import get_fex
 from qualitylib.tools import import_python_file, read_dataset
 from qualitylib.runner import Runner
 
@@ -9,17 +11,17 @@ def get_parser() -> argparse.ArgumentParser:
     parser.add_argument('--dataset', help='Path to dataset file for which to extract features', type=str)
     parser.add_argument('--fex_name', help='Name of feature extractor', type=str)
     parser.add_argument('--fex_version', help='Version of feature extractor', type=str, default=None)
-    parser.add_argument('--processes', help='Number of parallel processes', type=str, default=1)
+    parser.add_argument('--processes', help='Number of parallel processes', type=int, default=1)
     return parser
 
 
 def main() -> None:
     args = get_parser().parse_args()
-
     dataset = import_python_file(args.dataset)
     assets = read_dataset(dataset, shuffle=True)
-    runner = Runner(args.fex_name, args.fex_version, processes=args.processes, use_cache=True)  # Reads from stored results if available, else stores results.
-    runner(assets, return_results=False)  # Only extract features, do not use for anything.
+    FexClass = get_fex(args.fex_name, args.fex_version)  # Finds feature extractor by name and version
+    runner = Runner(FexClass, processes=args.processes, use_cache=True)  # Reads from stored results if available, else stores results.
+    runner(assets, return_results=False)  # Only extract features, do not use for anything. Since use_cache is True above, results are stored without being returned.
 
 
 if __name__ == '__main__':

diff --git a/examples/feature_extractors/ssim_fex.py b/examples/feature_extractors/ssim_fex.py
@@ -26,14 +26,15 @@ def _run_on_asset(self, asset_dict: Dict[str, Any]) -> Result:
                 standard=asset_dict['dis_standard'],
                 width=asset_dict['width'], height=asset_dict['height']
              ) as v_dis:
-                for i, (frame_ref, frame_dis) in enumerate(zip(v_ref, v_dis)):
-                    if i % sample_interval:
-                        continue
-                    lum_val, cs_val, ssim_val = ssim(frame_ref, frame_dis, full=True)
-                    feats_dict['lum'].append(lum_val)
-                    feats_dict['cs'].append(cs_val)
-                    feats_dict['ssim'].append(ssim_val)
+                # for i, (frame_ref, frame_dis) in enumerate(zip(v_ref, v_dis)):
+                #     if i % sample_interval:
+                #         continue
+                #     lum_val, cs_val, ssim_val = ssim(frame_ref, frame_dis, full=True)
+                #     feats_dict['lum'].append(lum_val)
+                #     feats_dict['cs'].append(cs_val)
+                #     feats_dict['ssim'].append(ssim_val)
+                feats = np.random.randn(v_ref.num_frames, 3)
 
-        feats = np.array(list(feats_dict.values())).T
+        # feats = np.array(list(feats_dict.values())).T
         print(f'Processed {asset_dict["dis_path"]}')
         return self._to_result(asset_dict, feats, list(feats_dict.keys()))