Allow old-style shape in blobproto_to_array

Fixes #3199
Bug introduced in #3170

python/caffe/io.py

@@ -20,10 +20,15 @@ def blobproto_to_array(blob, return_diff=False):
     Convert a blob proto to an array. In default, we will just return the data,
     unless return_diff is True, in which case we will return the diff.
     """
+    if blob.HasField('shape'):
+        shape = blob.shape.dim
+    else:
+        # Fall back to legacy 4D shape
+        shape = (blob.num, blob.channels, blob.height, blob.width)
     if return_diff:
-        return np.array(blob.diff).reshape(*blob.shape.dim)
+        return np.array(blob.diff).reshape(*shape)
     else:
-        return np.array(blob.data).reshape(*blob.shape.dim)
+        return np.array(blob.data).reshape(*shape)
 
 
 def array_to_blobproto(arr, diff=None):

python/caffe/test/test_io.py

@@ -0,0 +1,29 @@
+import numpy as np
+import unittest
+
+import caffe
+
+class TestBlobProtoToArray(unittest.TestCase):
+
+    def test_old_format(self):
+        data = np.zeros((10,10))
+        blob = caffe.proto.caffe_pb2.BlobProto()
+        blob.data.extend(list(data.flatten()))
+        shape = (1,1,10,10)
+        blob.num, blob.channels, blob.height, blob.width = shape
+
+        arr = caffe.io.blobproto_to_array(blob)
+        self.assertEqual(
+                arr.shape,
+                shape)
+
+    def test_new_format(self):
+        data = np.zeros((10,10))
+        blob = caffe.proto.caffe_pb2.BlobProto()
+        blob.data.extend(list(data.flatten()))
+        blob.shape.dim.extend(list(data.shape))
+
+        arr = caffe.io.blobproto_to_array(blob)
+        self.assertEqual(
+                arr.shape,
+                data.shape)