File size: 4,258 Bytes
1327f34 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | # Copyright 2025 The Scenic Authors.
#
# 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.
"""Implementation of data preprocessing ops for VTAB.
All preprocessing ops should return a data processing functors. A data
is represented as a dictionary of tensors, where field "image" is reserved
for 3D images (height x width x channels). The functors output dictionary with
field "image" being modified. Potentially, other fields can also be modified
or added.
"""
import numpy as np
from scenic.dataset_lib.big_transfer.registry import Registry
import tensorflow.compat.v1 as tf
@Registry.register("preprocess_ops.dsprites_pp", "function")
def get_dsprites_pp(predicted_attribute, num_classes=None):
"""Data preprocess function for dsprites dataset."""
attribute_to_classes = {
"label_shape": 3,
"label_scale": 6,
"label_orientation": 40,
"label_x_position": 32,
"label_y_position": 32,
}
def _dsprites_pp(data):
# For consistency with other datasets, image needs to have three channels
# and be in [0, 255). # pylint: disable=unused-argument
# data["image"] = tf.tile(data["image"], [1, 1, 3]) * 255
data["image"] = data["image"] * 255
# If num_classes is set, we group together nearby integer values to arrive
# at the desired number of classes. This is useful for example for grouping
# together different spatial positions.
num_original_classes = attribute_to_classes[predicted_attribute]
n_cls = num_original_classes if num_classes is None else num_classes
if not isinstance(n_cls, int) or n_cls <= 1 or n_cls > num_original_classes:
raise ValueError(
"The number of classes should be None or in [2, ..., num_classes].")
class_division_factor = float(num_original_classes) / n_cls
data["label"] = tf.cast(
tf.math.floordiv(
tf.cast(data[predicted_attribute], tf.float32),
class_division_factor), data[predicted_attribute].dtype)
return data
return _dsprites_pp
@Registry.register("preprocess_ops.clevr_pp", "function")
def get_clevr_pp(task, outkey="label"):
"""Data preprocess function for clevr dataset."""
def _count_preprocess_fn(data):
data[outkey] = tf.size(data["objects"]["size"]) - 3
return data
def _closest_object_preprocess_fn(data):
dist = tf.reduce_min(data["objects"]["pixel_coords"][:, 2])
# These thresholds are uniformly spaced and result in more or less balanced
# distribution of classes.
thrs = np.array([0.0, 8.0, 8.5, 9.0, 9.5, 10.0, 100.0])
data[outkey] = tf.reduce_max(tf.where((thrs - dist) < 0))
return data
task_to_preprocess = {
"count_all": _count_preprocess_fn,
"closest_object_distance": _closest_object_preprocess_fn,
}
return task_to_preprocess[task]
@Registry.register("preprocess_ops.kitti_pp", "function")
def get_kitti_pp(task):
"""Data preprocess function for kitti dataset."""
def _closest_vehicle_distance_pp(data):
"""Predict the distance to the closest vehicle."""
# Location feature contains (x, y, z) in meters w.r.t. the camera.
vehicles = tf.where(data["objects"]["type"] < 3) # Car, Van, Truck.
vehicle_z = tf.gather(
params=data["objects"]["location"][:, 2], indices=vehicles)
vehicle_z = tf.concat([vehicle_z, tf.constant([[1000.0]])], axis=0)
dist = tf.reduce_min(vehicle_z)
# Results in a uniform distribution over three distances, plus one class for
# "no vehicle".
thrs = np.array([-100.0, 8.0, 20.0, 999.0])
label = tf.reduce_max(tf.where((thrs - dist) < 0))
return {"image": data["image"], "label": label}
task_to_preprocess = {
"closest_vehicle_distance": _closest_vehicle_distance_pp,
}
return task_to_preprocess[task]
|