refactor(preprocessor): separate resize logic from the unary preproce…

…ssor

gnes/base/__init__.py

@@ -28,7 +28,7 @@
 
 from ..helper import set_logger, profiling, yaml, parse_arg, load_contrib_module
 
-__all__ = ['TrainableBase']
+__all__ = ['TrainableBase', 'CompositionalTrainableBase']
 
 T = TypeVar('T', bound='TrainableBase')
 

gnes/preprocessor/__init__.py

@@ -29,6 +29,7 @@
     'WeightedSlidingPreprocessor': 'image.sliding_window',
     'SegmentPreprocessor': 'image.segmentation',
     'BaseUnaryPreprocessor': 'base',
+    'ResizeChunkPreprocessor': 'image.resize',
     'BaseVideoPreprocessor': 'video.base',
     'FFmpegPreprocessor': 'video.ffmpeg',
     'ShotDetectPreprocessor': 'video.shotdetect',

gnes/preprocessor/base.py

@@ -28,14 +28,21 @@
 class BasePreprocessor(TrainableBase):
     doc_type = gnes_pb2.Document.UNKNOWN
 
-    def __init__(self, start_doc_id: int = 0, random_doc_id: bool = True, *args, **kwargs):
+    def __init__(self, start_doc_id: int = 0,
+                 random_doc_id: bool = True,
+                 uniform_doc_weight: bool = True,
+                 *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.start_doc_id = start_doc_id
         self.random_doc_id = random_doc_id
+        self.uniform_doc_weight = uniform_doc_weight
 
     def apply(self, doc: 'gnes_pb2.Document') -> None:
         doc.doc_id = self.start_doc_id if not self.random_doc_id else random.randint(0, ctypes.c_uint(-1).value)
         doc.doc_type = self.doc_type
+        if not doc.weight and self.uniform_doc_weight:
+            doc.weight = 1.0
+        self.start_doc_id += 1
 
 
 class PipelinePreprocessor(CompositionalTrainableBase):
@@ -55,11 +62,10 @@ def train(self, data, *args, **kwargs):
 
 
 class BaseUnaryPreprocessor(BasePreprocessor):
+    is_trained = True
 
     def __init__(self, doc_type: int, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.target_img_size = 224
-        self.is_trained = True
         self.doc_type = doc_type
 
     def apply(self, doc: 'gnes_pb2.Document'):
@@ -78,7 +84,6 @@ def raw_to_chunk(self, chunk: 'gnes_pb2.Chunk', raw_bytes: bytes):
             chunk.text = raw_bytes.decode()
         elif self.doc_type == gnes_pb2.Document.IMAGE:
             img = np.array(Image.open(io.BytesIO(raw_bytes)))
-            img = np.array(Image.fromarray(img).resize((self.target_img_size, self.target_img_size)))
             chunk.blob.CopyFrom(array2blob(img))
         elif self.doc_type == gnes_pb2.Document.VIDEO:
             raise NotImplementedError

gnes/preprocessor/image/base.py

@@ -46,4 +46,4 @@ def _get_all_subarea(image):
 
         index = [[x, y, x + 1, y + 1] for [x, y] in product(range(len(x_list) - 1), range(len(y_list) - 1))]
         all_subareas = [[x_list[idx[0]], y_list[idx[1]], x_list[idx[2]], y_list[idx[3]]] for idx in index]
-        return all_subareas, index
+        return all_subareas, index

gnes/preprocessor/image/resize.py

@@ -0,0 +1,38 @@
+#  Tencent is pleased to support the open source community by making GNES available.
+#
+#  Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import numpy as np
+from PIL import Image
+
+from .base import BaseImagePreprocessor
+from ...proto import gnes_pb2, blob2array, array2blob
+
+
+class ResizeChunkPreprocessor(BaseImagePreprocessor):
+
+    def __init__(self,
+                 target_width: int = 224,
+                 target_height: int = 224,
+                 *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.target_width = target_width
+        self.target_height = target_height
+
+    def apply(self, doc: 'gnes_pb2.Document') -> None:
+        super().apply(doc)
+        for c in doc.chunks:
+            img = blob2array(c.blob)
+            img = np.array(Image.fromarray(img).resize((self.target_width, self.target_height)))
+            c.blob.CopyFrom(array2blob(img))

gnes/preprocessor/image/segmentation.py

@@ -28,13 +28,11 @@ class SegmentPreprocessor(BaseImagePreprocessor):
 
     def __init__(self, model_name: str,
                  model_dir: str,
-                 target_img_size: int = 224,
                  _use_cuda: bool = False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.model_name = model_name
         self.model_dir = model_dir
-        self.target_img_size = target_img_size
         self._use_cuda = _use_cuda
 
     def post_init(self):
@@ -68,7 +66,7 @@ def apply(self, doc: 'gnes_pb2.Document'):
             for ci, ele in enumerate(zip(chunks, weight)):
                 c = doc.chunks.add()
                 c.doc_id = doc.doc_id
-                c.blob.CopyFrom(array2blob(self._crop_image_reshape(original_image, ele[0])))
+                c.blob.CopyFrom(array2blob(self._crop_resize(original_image, ele[0])))
                 c.offset_1d = ci
                 c.offset_nd.x.extend(self._get_seg_offset_nd(all_subareas, index, ele[0]))
                 c.weight = self._cal_area(ele[0]) / (original_image.size[0] * original_image.size[1])
@@ -83,7 +81,7 @@ def apply(self, doc: 'gnes_pb2.Document'):
         else:
             self.logger.error('bad document: "raw_bytes" is empty!')
 
-    def _crop_image_reshape(self, original_image, coordinates):
+    def _crop_resize(self, original_image, coordinates):
         return np.array(original_image.crop(coordinates).resize((self.target_img_size,
                                                                  self.target_img_size)))
 

gnes/preprocessor/image/sliding_window.py

@@ -29,13 +29,11 @@ class BaseSlidingPreprocessor(BaseImagePreprocessor):
     def __init__(self, window_size: int = 64,
                  stride_height: int = 64,
                  stride_wide: int = 64,
-                 target_img_size: int = 224,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.window_size = window_size
         self.stride_height = stride_height
         self.stride_wide = stride_wide
-        self.target_img_size = target_img_size
 
     def apply(self, doc: 'gnes_pb2.Document'):
         super().apply(doc)

tests/test_image_preprocessor.py

@@ -15,6 +15,7 @@ def setUp(self):
         self.unary_img_pre_yaml = os.path.join(self.dirname, 'yaml', 'base-unary-image-prep.yml')
         self.slidingwindow_img_pre_yaml = os.path.join(self.dirname, 'yaml', 'base-vanilla_sldwin-image-prep.yml')
         self.segmentation_img_pre_yaml = os.path.join(self.dirname, 'yaml', 'base-segmentation-image-prep.yml')
+        self.resize_img_pre_yaml = os.path.join(self.dirname, 'yaml', 'resize-image-prep.yml')
 
     def test_unary_preprocessor_service_empty(self):
         args = set_preprocessor_service_parser().parse_args([
@@ -116,9 +117,31 @@ def test_unary_preprocessor_service_realdata(self):
                     self.assertEqual(len(d.chunks), 1)
                     self.assertEqual(len(blob2array(d.chunks[0].blob).shape), 3)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[-1], 3)
+
+    def test_resize_preprocessor_service_realdata(self):
+        args = set_preprocessor_service_parser().parse_args([
+            '--yaml_path', self.resize_img_pre_yaml
+        ])
+        c_args = _set_client_parser().parse_args([
+            '--port_in', str(args.port_out),
+            '--port_out', str(args.port_in)
+        ])
+        all_zips = zipfile.ZipFile(os.path.join(self.dirname, 'imgs/test.zip'))
+        all_bytes = [all_zips.open(v).read() for v in all_zips.namelist()]
+
+        with PreprocessorService(args), ZmqClient(c_args) as client:
+            for req in RequestGenerator.index(all_bytes):
+                msg = gnes_pb2.Message()
+                msg.request.index.CopyFrom(req.index)
+                client.send_message(msg)
+                r = client.recv_message()
+                self.assertEqual(r.envelope.routes[0].service, 'PreprocessorService:PipelinePreprocessor')
+                for d in r.request.index.docs:
+                    self.assertEqual(len(d.chunks), 1)
+                    self.assertEqual(len(blob2array(d.chunks[0].blob).shape), 3)
+                    self.assertEqual(blob2array(d.chunks[0].blob).shape[-1], 3)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[0], 224)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[1], 224)
-                    print(blob2array(d.chunks[0].blob).dtype)
 
     def test_slidingwindow_preprocessor_service_realdata(self):
         args = set_preprocessor_service_parser().parse_args([
@@ -144,7 +167,6 @@ def test_slidingwindow_preprocessor_service_realdata(self):
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[-1], 3)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[0], 224)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[1], 224)
-                    print(blob2array(d.chunks[0].blob).dtype)
 
     def test_segmentation_preprocessor_service_realdata(self):
         args = set_preprocessor_service_parser().parse_args([
@@ -170,4 +192,4 @@ def test_segmentation_preprocessor_service_realdata(self):
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[-1], 3)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[0], 224)
                     self.assertEqual(blob2array(d.chunks[0].blob).shape[1], 224)
-                    print(blob2array(d.chunks[0].blob).dtype)
+                    print(blob2array(d.chunks[0].blob).dtype)

tests/yaml/resize-image-prep.yml

@@ -0,0 +1,9 @@
+!PipelinePreprocessor
+component:
+  - !BaseUnaryPreprocessor
+    parameter:
+      doc_type: 2
+  - !ResizeChunkPreprocessor
+    parameter:
+      target_height: 224
+      target_width: 224