refactor

tools/benchmark.lua

@@ -42,7 +42,7 @@ local function rgb2y_matlab(x)
    return y:byte():float()
 end
 
-local function MSE(x1, x2)
+local function RGBMSE(x1, x2)
    x1 = iproc.float2byte(x1):float()
    x2 = iproc.float2byte(x2):float()
    return (x1 - x2):pow(2):mean()
@@ -58,15 +58,17 @@ local function YMSE(x1, x2)
       return (x1_2 - x2_2):pow(2):mean()
    end
 end
-local function PSNR(x1, x2)
-   local mse = MSE(x1, x2)
-   return 10 * math.log10((255.0 * 255.0) / mse)
+local function MSE(x1, x2, color)
+   if color == "y" then
+      return YMSE(x1, x2)
+   else
+      return RGBMSE(x1, x2)
+   end
 end
-local function YPSNR(x1, x2)
-   local mse = YMSE(x1, x2)
+local function PSNR(x1, x2, color)
+   local mse = MSE(x1, x2, color)
    return 10 * math.log10((255.0 * 255.0) / mse)
 end
-
 local function transform_jpeg(x, opt)
    for i = 1, opt.jpeg_times do
       jpeg = gm.Image(x, "RGB", "DHW")
@@ -117,30 +119,15 @@ local function benchmark(opt, x, input_func, model1, model2)
 	 end
 	 baseline_output = baseline_scale(input, opt.filter)
       end
-      if opt.color == "y" then
-	 model1_mse = model1_mse + YMSE(ground_truth, model1_output)
-	 model1_psnr = model1_psnr + YPSNR(ground_truth, model1_output)
-	 if model2 then
-	    model2_mse = model2_mse + YMSE(ground_truth, model2_output)
-	    model2_psnr = model2_psnr + YPSNR(ground_truth, model2_output)
-	 end
-	 if baseline_output then
-	    baseline_mse = baseline_mse + YMSE(ground_truth, baseline_output)
-	    baseline_psnr = baseline_psnr + YPSNR(ground_truth, baseline_output)
-	 end
-      elseif opt.color == "rgb" then
-	 model1_mse = model1_mse + MSE(ground_truth, model1_output)
-	 model1_psnr = model1_psnr + PSNR(ground_truth, model1_output)
-	 if model2 then
-	    model2_mse = model2_mse + MSE(ground_truth, model2_output)
-	    model2_psnr = model2_psnr + PSNR(ground_truth, model2_output)
-	 end
-	 if baseline_output then
-	    baseline_mse = baseline_mse + MSE(ground_truth, baseline_output)
-	    baseline_psnr = baseline_psnr + PSNR(ground_truth, baseline_output)
-	 end
-      else
-	 error("Unknown color: " .. opt.color)
+      model1_mse = model1_mse + MSE(ground_truth, model1_output, opt.color)
+      model1_psnr = model1_psnr + PSNR(ground_truth, model1_output, opt.color)
+      if model2 then
+	 model2_mse = model2_mse + MSE(ground_truth, model2_output, opt.color)
+	 model2_psnr = model2_psnr + PSNR(ground_truth, model2_output, opt.color)
+      end
+      if baseline_output then
+	 baseline_mse = baseline_mse + MSE(ground_truth, baseline_output, opt.color)
+	 baseline_psnr = baseline_psnr + PSNR(ground_truth, baseline_output, opt.color)
       end
       if model2 then
 	 if baseline_output then