google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""utils for writing video."""
17
import base64
18
import io
19
import re
20
import PIL
21
import skvideo.io
22

23

24
def get_writer(file_name, frame_rate=30, pix_fmt='yuv420p', vcodec='libx264'):
25
  rate = str(frame_rate)
26
  inputdict = {'-r': rate}
27
  outputdict = {
28
      '-pix_fmt': pix_fmt,
29
      '-r': rate,
30
      '-vcodec': vcodec,
31
  }
32
  writer = skvideo.io.FFmpegWriter(file_name, inputdict, outputdict)
33
  return writer
34

35

36
def save_image(url, path):
37
  data = re.sub('^data:image/.+;base64,', '', url)
38
  image = PIL.Image.open(io.BytesIO(base64.b64decode(data)))
39
  image.save(path)
40