-a predict with RGB image converts to Gray

Laulian · Dec 22, 2015 · 834f9a8 · 834f9a8
1 parent fb205e0
commit 834f9a8
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/main.lua b/main.lua
@@ -1077,8 +1077,14 @@ function stereo_predict(x_batch, id)
 end
 
 if opt.a == 'predict' then
-   x0 = image.load(opt.left, 1, 'byte'):float()
-   x1 = image.load(opt.right, 1, 'byte'):float()
+   x0 = image.load(opt.left, nil, 'byte'):float()
+   x1 = image.load(opt.right, nil, 'byte'):float()
+
+   if x0:size(1) == 3 then
+      assert(x1:size(1) == 3)
+      x0 = image.rgb2y(x0)
+      x1 = image.rgb2y(x1)
+   end
    disp_max = opt.disp_max
 
    x0:add(-x0:mean()):div(x0:std())