File size: 2,971 Bytes
74e8f2f |
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 |
# Copyright 2024 Big Vision 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.
"""Tests for ops_image."""
import copy
import io
import big_vision.pp.ops_image as pp
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
def get_image_data():
img = tf.random.uniform((640, 480, 3), 0, 255, tf.int32) # Can't ask uint8!?
return {"image": tf.cast(img, tf.uint8)}
class PreprocessOpsTest(tf.test.TestCase):
def tfrun(self, ppfn, data):
# Run once as standalone, as could happen eg in colab.
yield {k: np.array(v) for k, v in ppfn(copy.deepcopy(data)).items()}
# And then once again as part of tfdata pipeline.
# You'd be surprised how much these two differ!
tfdata = tf.data.Dataset.from_tensors(copy.deepcopy(data))
for npdata in tfdata.map(ppfn).as_numpy_iterator():
yield npdata
def test_resize(self):
for data in self.tfrun(pp.get_resize([120, 80]), get_image_data()):
self.assertEqual(data["image"].shape, (120, 80, 3))
def test_resize_small(self):
for data in self.tfrun(pp.get_resize_small(240), get_image_data()):
self.assertEqual(data["image"].shape, (320, 240, 3))
def test_resize_long(self):
for data in self.tfrun(pp.get_resize_long(320), get_image_data()):
self.assertEqual(data["image"].shape, (320, 240, 3))
def test_inception_crop(self):
for data in self.tfrun(pp.get_inception_crop(), get_image_data()):
self.assertEqual(data["image"].shape[-1], 3)
def test_decode_jpeg_and_inception_crop(self):
f = io.BytesIO()
plt.imsave(f, get_image_data()["image"].numpy(), format="jpg")
data = {"image": tf.cast(f.getvalue(), tf.string)}
for data in self.tfrun(pp.get_decode_jpeg_and_inception_crop(), data):
self.assertEqual(data["image"].shape[-1], 3)
def test_random_crop(self):
for data in self.tfrun(pp.get_random_crop([120, 80]), get_image_data()):
self.assertEqual(data["image"].shape, (120, 80, 3))
def test_central_crop(self):
for data in self.tfrun(pp.get_central_crop([20, 80]), get_image_data()):
self.assertEqual(data["image"].shape, (20, 80, 3))
def test_random_flip_lr(self):
data_orig = get_image_data()
for data in self.tfrun(pp.get_random_flip_lr(), data_orig):
self.assertTrue(
np.all(data_orig["image"].numpy() == data["image"]) or
np.all(data_orig["image"].numpy() == data["image"][:, ::-1]))
if __name__ == "__main__":
tf.test.main()
|