File size: 2,473 Bytes
b6ff324 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | # ------------------------------------------------------------------------
# Copyright (c) 2024-present, BAAI. 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.
# ------------------------------------------------------------------------
"""Flex data transforms."""
import re
import numpy as np
import numpy.random as npr
class Transform(object):
"""Base transform type."""
def filter_outputs(self, *outputs):
outputs = [x for x in outputs if x is not None]
return outputs if len(outputs) > 1 else outputs[0]
class ParseLatents(Transform):
"""Parse VQ or VAE latents."""
def __init__(self):
super().__init__()
def __call__(self, inputs):
for k, dtype in zip(("moments", "codes"), ("float16", "int32")):
if k in inputs:
return np.frombuffer(inputs[k], dtype).reshape(inputs["shape"])
raise ValueError("Missing latents in inputs.")
class ParseAnnotations(Transform):
"""Parse ground-truth annotations."""
def __init__(self, short_prob=0.5):
super().__init__()
self.short_prob = short_prob
def __call__(self, inputs):
text = inputs.get("text", None)
label = inputs.get("label", None)
caption = inputs.get("caption", None)
if caption and isinstance(caption, dict): # Cached.
caption = np.frombuffer(caption["data"], "float16").reshape(caption["shape"])
if text and isinstance(text, dict) and len(text["data"]) > 0 and npr.rand() < 0.5:
caption = np.frombuffer(text["data"], "float16").reshape(text["shape"])
return label, caption
# Improved short caption.
if label is None:
text_match = re.match(r"^(.*?[.!?])\s+", caption)
text = text if text else (text_match.group(1) if text_match else caption)
caption = text if text and npr.rand() < self.short_prob else caption
return label, caption
|