Merge pull request #286 from jina-ai/fix-drivers-blob-285

fix(drivers): add handling for blob

jina/drivers/craft.py

@@ -2,7 +2,7 @@
 import random
 
 from . import BaseExecutableDriver
-from .helper import array2blob, pb_obj2dict
+from .helper import array2blob, pb_obj2dict, blob2array
 
 
 class BaseCraftDriver(BaseExecutableDriver):
@@ -26,10 +26,16 @@ def __call__(self, *args, **kwargs):
                 continue
             _chunks_to_add = []
             for c in d.chunks:
-                ret = self.exec_fn(**pb_obj2dict(c, self.exec.required_keys))
+                _args_dict = pb_obj2dict(c, self.exec.required_keys)
+                if 'blob' in self.exec.required_keys:
+                    _args_dict['blob'] = blob2array(c.blob)
+                ret = self.exec_fn(**_args_dict)
                 if isinstance(ret, dict):
                     for k, v in ret.items():
-                        setattr(c, k, v)
+                        if k == 'blob':
+                            c.blob.CopyFrom(array2blob(v))
+                        else:
+                            setattr(c, k, v)
                     continue
                 if isinstance(ret, list):
                     for chunk_dict in ret:

jina/executors/indexers/vector/numpy.py

@@ -132,7 +132,6 @@ def query(self, keys: np.ndarray, top_k: int, *args, **kwargs) -> Tuple['np.ndar
 
         idx = dist.argsort(axis=1)[:, :top_k]
         dist = np.take_along_axis(dist, idx, axis=1)
-
         return self.int2ext_key[idx], dist
 
 

tests/executors/encoders/image/test_torchvision.py

@@ -7,6 +7,7 @@
 class MyTestCase(ImageTestCase):
     def _get_encoder(self):
         self.target_output_dim = 1280
+        self.input_dim = 224
         return ImageTorchEncoder()