Merge pull request #325 from gnes-ai/webp_encoder

feat(webp-encoder): add webp encoder

gnes/client/base.py

@@ -78,6 +78,8 @@ def __init__(self, args):
         self.logger = set_logger(self.__class__.__name__, self.args.verbose)
         self.ctx = zmq.Context()
         self.ctx.setsockopt(zmq.LINGER, 0)
+        self.logger.info('current libzmq version is %s,  pyzmq version is %s' % (zmq.zmq_version(), zmq.__version__))
+
         self.receiver, recv_addr = build_socket(
             self.ctx, self.args.host_in, self.args.port_in,
             self.args.socket_in, self.args.identity)

gnes/preprocessor/io_utils/webp.py

@@ -0,0 +1,35 @@
+#  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.
+
+from typing import List
+
+def encode_video(images: List['np.ndarray'], frame_rate: int, pix_fmt: str = 'rgb24'):
+    import webp
+    from PIL import Image
+
+    height, width, _ = images[0].shape
+    if pix_fmt == 'rgb24':
+        pix_fmt = 'RGB'
+    # Save an animation
+    enc = webp.WebPAnimEncoder.new(width, height)
+    timestamp_ms = 0
+    duration = 1000 // frame_rate
+    for x in images:
+        img = Image.fromarray(x.copy(), pix_fmt)
+        pic = webp.WebPPicture.from_pil(img)
+        enc.encode_frame(pic, timestamp_ms)
+        timestamp_ms += duration
+    anim_data = enc.assemble(timestamp_ms)
+    return bytes(anim_data.buffer())

gnes/preprocessor/video/video_encoder.py

@@ -15,7 +15,7 @@
 
 from ...proto import gnes_pb2, blob2array
 from ..base import BaseVideoPreprocessor
-from ..io_utils import video, gif
+from ..io_utils import video, gif, webp
 
 
 class VideoEncoderPreprocessor(BaseVideoPreprocessor):
@@ -25,16 +25,17 @@ def __init__(self, frame_rate: int = 10, pix_fmt: str = 'rgb24', video_format: s
         self.frame_rate = frame_rate
         self.video_format = video_format
 
-        if self.video_format not in ['mp4', 'gif']:
+        if self.video_format not in ['mp4', 'gif', 'webp']:
             raise ValueError("%s encoder has not been supported!" % (self.video_format))
 
-
     def _encode(self, images: 'np.ndarray'):
         encoder = None
         if self.video_format == 'mp4':
             encoder = video
         elif self.video_format == 'gif':
             encoder = gif
+        elif self.video_format == 'webp':
+            encoder = webp
 
         return encoder.encode_video(images, pix_fmt=self.pix_fmt, frame_rate=self.frame_rate)
 

gnes/service/frontend.py

@@ -120,6 +120,8 @@ def StreamCall(self, request_iterator, context):
             self.pending_request = 0
 
             def get_response(num_recv, blocked=False):
+                if blocked:
+                    self.logger.info('waiting for %d responses ...' % (num_recv))
                 for _ in range(num_recv):
                     if blocked or zmq_client.receiver.poll(1):
                         msg = zmq_client.recv_message(**self.send_recv_kwargs)
@@ -129,12 +131,14 @@ def get_response(num_recv, blocked=False):
             with self.zmq_context as zmq_client:
 
                 for request in request_iterator:
+                    self.logger.info('receive request: %s' % request.request_id)
                     num_recv = max(self.pending_request - self.args.max_pending_request, 1)
                     yield from get_response(num_recv, num_recv > 1)
-
+                    self.logger.info('send new request into %d appending tasks' % (self.pending_request))
                     zmq_client.send_message(self.add_envelope(request, zmq_client), **self.send_recv_kwargs)
                     self.pending_request += 1
 
+                self.logger.info('all requests are sent, waiting for the responses...')
                 yield from get_response(self.pending_request, blocked=True)
 
         class ZmqContext: