|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
"""Tests for preprocessing utils.""" |
|
|
|
from __future__ import absolute_import |
|
from __future__ import division |
|
from __future__ import print_function |
|
|
|
from big_vision.pp import utils |
|
import tensorflow.compat.v1 as tf |
|
|
|
|
|
class UtilsTest(tf.test.TestCase): |
|
|
|
def test_maybe_repeat(self): |
|
self.assertEqual((1, 1, 1), utils.maybe_repeat(1, 3)) |
|
self.assertEqual((1, 2), utils.maybe_repeat((1, 2), 2)) |
|
self.assertEqual([1, 2], utils.maybe_repeat([1, 2], 2)) |
|
|
|
def test_inkeyoutkey(self): |
|
@utils.InKeyOutKey() |
|
def get_pp_fn(shift, scale=0): |
|
def _pp_fn(x): |
|
return scale * x + shift |
|
return _pp_fn |
|
|
|
data = {"k_in": 2, "other": 3} |
|
ppfn = get_pp_fn(1, 2, inkey="k_in", outkey="k_out") |
|
self.assertEqual({"k_in": 2, "k_out": 5, "other": 3}, ppfn(data)) |
|
|
|
data = {"k": 6, "other": 3} |
|
ppfn = get_pp_fn(1, inkey="k", outkey="k") |
|
self.assertEqual({"k": 1, "other": 3}, ppfn(data)) |
|
|
|
data = {"other": 6, "image": 3} |
|
ppfn = get_pp_fn(5, 2) |
|
self.assertEqual({"other": 6, "image": 11}, ppfn(data)) |
|
|
|
|
|
if __name__ == "__main__": |
|
tf.test.main() |
|
|