Back to main code from test version

abhinaukumar · Oct 31, 2023 · eb42c1e · eb42c1e
1 parent b0cd08c
commit eb42c1e
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/examples/feature_extractors/ssim_fex.py b/examples/feature_extractors/ssim_fex.py
@@ -26,15 +26,14 @@ 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)
-                feats = np.random.randn(v_ref.num_frames, 3)
+                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.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()))