feat(#39): add sim chart

- ssim - mse - psnr
williamfzc · Aug 23, 2019 · 78e8ec6 · 78e8ec6
1 parent d96dbf2
commit 78e8ec6
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/stagesepx/reporter.py b/stagesepx/reporter.py
@@ -9,6 +9,7 @@
 
 from stagesepx.classifier import ClassifierResult
 from stagesepx import toolbox
+from stagesepx.cutter import VideoCutResult
 
 TEMPLATE = r'''
 <!DOCTYPE html>
@@ -166,6 +167,24 @@ def _draw_line(data_list: typing.List[ClassifierResult]) -> Line:
         )
         return line
 
+    @staticmethod
+    def _draw_sim(data: VideoCutResult) -> Line:
+        x_axis = [str(i.start) for i in data.range_list]
+        ssim_axis = [i.ssim for i in data.range_list]
+        mse_axis = [i.mse for i in data.range_list]
+
+        line = Line()
+        line.add_xaxis(x_axis)
+        line.add_yaxis('ssim', ssim_axis)
+        line.add_yaxis('mse', mse_axis)
+        line.set_global_opts(
+            title_opts=opts.TitleOpts(title='SIM'),
+            toolbox_opts=opts.ToolboxOpts(is_show=True),
+            tooltip_opts=opts.TooltipOpts(is_show=True, trigger='axis', axis_pointer_type='cross'),
+        )
+        line.set_series_opts(label_opts=opts.LabelOpts(is_show=False))
+        return line
+
     @staticmethod
     def _draw_bar(data_list: typing.List[ClassifierResult]) -> Bar:
         # draw bar chart
@@ -189,12 +208,14 @@ def _draw_bar(data_list: typing.List[ClassifierResult]) -> Bar:
 
     def draw(self,
              data_list: typing.List[ClassifierResult],
-             report_path: str = None):
+             report_path: str = None,
+             cut_result: VideoCutResult = None):
         """
         draw report file
 
         :param data_list: classifierResult list, output of classifier
         :param report_path: your report will be there
+        :param cut_result: more charts would be built
         :return:
         """
 
@@ -207,6 +228,10 @@ def draw(self,
         page.add(line)
         page.add(bar)
 
+        if cut_result:
+            sim_line = self._draw_sim(cut_result)
+            page.add(sim_line)
+
         # insert extras
         template = Template(TEMPLATE)
         template_content = template.render(

diff --git a/stagesepx/toolbox.py b/stagesepx/toolbox.py
@@ -6,6 +6,7 @@
 import numpy as np
 from skimage.filters import threshold_otsu
 from skimage.measure import compare_ssim as origin_compare_ssim
+from skimage.measure import compare_nrmse, compare_psnr
 from skimage.feature import hog, local_binary_pattern
 
 
@@ -161,7 +162,13 @@ def sharpen_frame(old: np.ndarray) -> np.ndarray:
 
 def calc_mse(pic1: np.ndarray, pic2: np.ndarray) -> float:
     # MSE: https://en.wikipedia.org/wiki/Mean_squared_error
-    return np.sum((pic1.astype('float') - pic2.astype('float')) ** 2) / float(pic1.shape[0] * pic2.shape[1])
+    # return np.sum((pic1.astype('float') - pic2.astype('float')) ** 2) / float(pic1.shape[0] * pic2.shape[1])
+    return compare_nrmse(pic1, pic2)
+
+
+def calc_psnr(pic1: np.ndarray, pic2: np.ndarray) -> float:
+    # PSNR: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
+    return compare_psnr(pic1, pic2)
 
 
 def compress_frame(old: np.ndarray,