Spaces:
Sleeping
Sleeping
File size: 2,353 Bytes
5b21b68 | 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 | import json
import tempfile
import unittest
import numpy as np
from PIL import Image
class FakePolicy:
def __init__(self, actions=None):
self.observation = None
self.noise = None
self.actions = np.ones((10, 7)) if actions is None else actions
def infer(self, observation, *, noise=None):
self.observation = observation
self.noise = noise
return {"actions": self.actions}
class InferenceTests(unittest.TestCase):
def test_prediction_maps_ur_observation_and_writes_json(self):
from inference import ACTION_LABELS, run_prediction
policy = FakePolicy()
image = Image.new("RGB", (8, 6), "red")
result = run_prediction(
policy, image, image, "pick up", [1, 2, 3, 4, 5, 6, 0], 2,
"owner/model", "checkpoint",
)
self.assertEqual(tuple(result.actions.columns), ACTION_LABELS)
self.assertEqual(
policy.observation["observation/state"].tolist(),
[1, 2, 3, 4, 5, 6, 0],
)
self.assertEqual(policy.observation["observation/image"].dtype, np.uint8)
self.assertEqual(policy.observation["prompt"], "pick up")
self.assertEqual(policy.noise.shape, (10, 32))
with open(result.json_path, encoding="utf-8") as stream:
self.assertEqual(json.load(stream)["seed"], 44)
def test_same_trial_uses_same_diffusion_noise(self):
from inference import run_prediction
first, second = FakePolicy(), FakePolicy()
image = Image.new("RGB", (8, 6))
arguments = (image, image, "task", [0] * 7, 3, "model", "checkpoint")
run_prediction(first, *arguments)
run_prediction(second, *arguments)
np.testing.assert_array_equal(first.noise, second.noise)
def test_prediction_rejects_invalid_state_and_action_shape(self):
from inference import run_prediction
image = Image.new("RGB", (8, 6))
with self.assertRaisesRegex(ValueError, "seven"):
run_prediction(FakePolicy(), image, image, "task", [1, 2], 0, "m", "c")
with self.assertRaisesRegex(RuntimeError, r"\(10, 7\)"):
run_prediction(
FakePolicy(np.zeros((8, 7))), image, image, "task", [0] * 7,
0, "m", "c",
)
if __name__ == "__main__":
unittest.main()
|