follow-up on #443 to invert k channels (instead of 3)

BVLC · May 23, 2014 · 211c3c2 · 211c3c2
1 parent e5fcf03
commit 211c3c2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/caffe/pycaffe.py b/python/caffe/pycaffe.py
@@ -284,7 +284,8 @@ def _Net_deprocess(self, input_name, input_):
         decaf_in += mean
     decaf_in = decaf_in.transpose((1,2,0))
     if channel_order:
-        channel_order_inverse = tuple([channel_order.index(i) for i in (0,1,2)])
+        channel_order_inverse = [channel_order.index(i)
+                                 for i in range(decaf_in.shape[2])]
         decaf_in = decaf_in[:, :, channel_order_inverse]
     if input_scale:
         decaf_in /= input_scale