def test_repeated_max_len():
repeated = Repeated(dtype=Label(classes=["classA", "classB"]), max_len=2)
with pytest.raises(ValueError):
repeated.deserialize(*({
"label": "classA"
}, {
"label": "classA"
}, {
"label": "classB"
}))
assert repeated.deserialize(*({
"label": "classA"
}, {
"label": "classB"
})) == (
torch.tensor(0),
torch.tensor(1),
)
with pytest.raises(ValueError):
repeated.serialize((torch.tensor(0), torch.tensor(0), torch.tensor(1)))
assert repeated.serialize(
(torch.tensor(1), torch.tensor(0))) == ("classB", "classA")
# max_len < 1
with pytest.raises(ValueError):
Repeated(dtype=Label(classes=["classA", "classB"]), max_len=0)
assert Repeated(dtype=Label(classes=["classA", "classB"]),
max_len=1) is not None
# type(max_len) is not int
with pytest.raises(TypeError):
Repeated(dtype=Label(classes=["classA", "classB"]), max_len=str)
class ClassificationInference(ModelComponent):
def __init__(self, model): # skipcq: PYL-W0621
self.model = model
@expose(
inputs={"img": Image(extension="JPG")},
outputs={"prediction": Label(path=str(CWD / "imagenet_labels.txt"))},
)
def classify(self, img):
img = img.float() / 255
mean = torch.tensor([[[0.485, 0.456, 0.406]]]).float()
std = torch.tensor([[[0.229, 0.224, 0.225]]]).float()
img = (img - mean) / std
img = img.permute(0, 3, 2, 1)
out = self.model(img)
method_res = self.method_from_exposed(42)
assert method_res == ("method", 42)
func_res = _func_from_exposed("DouglasAdams")
assert func_res == ("func", "DouglasAdams")
return out.argmax()
@staticmethod
def never_should_run():
raise RuntimeError()
@staticmethod
def method_from_exposed(arg):
return ("method", arg)
def test_repeated_serialize(session_global_datadir):
repeated = Repeated(dtype=Label(path=str(session_global_datadir / "imagenet_labels.txt")))
assert repeated.deserialize(*({"label": "chickadee"}, {"label": "stingray"})) == (
torch.tensor(19),
torch.tensor(6),
)
assert repeated.serialize((torch.tensor(19), torch.tensor(6))) == ("chickadee", "stingray")
assert repeated.serialize(torch.tensor([19, 6])) == ("chickadee", "stingray")
def test_repeated_deserialize():
repeated = Repeated(dtype=Label(classes=["classA", "classB"]))
res = repeated.deserialize(*({
"label": "classA"
}, {
"label": "classA"
}, {
"label": "classB"
}))
assert res == (torch.tensor(0), torch.tensor(0), torch.tensor(1))
class SeatClassifier(ModelComponent):
def __init__(self, model, config):
self.sport = config["sport"]
@expose(
inputs={
"section": Number(),
"isle": Number(),
"row": Number(),
"stadium": Label(path=str(CWD / "imagenet_labels.txt")),
},
outputs={
"seat_number": Number(),
"team": Label(path=str(CWD / "imagenet_labels.txt")),
},
)
def predict(self, section, isle, row, stadium):
seat_num = section.item() * isle.item() * row.item(
) * stadium * len(self.sport)
stadium_idx = torch.tensor(1000)
return torch.Tensor([seat_num]), stadium_idx
class ClassificationInferenceComposable(ModelComponent):
def __init__(self, model):
self.model = model
@expose(
inputs={
"img": Image(extension="JPG"),
"tag": Label(path=str(CWD / "imagenet_labels.txt")),
},
outputs={
"predicted_tag": Label(path=str(CWD / "imagenet_labels.txt")),
"cropped_img": Image(),
},
)
def classify(self, img, tag):
im_div = img.float() / 255
mean = torch.tensor([[[0.485, 0.456, 0.406]]]).float()
std = torch.tensor([[[0.229, 0.224, 0.225]]]).float()
img_new = (im_div - torch.mean(mean)) / torch.mean(std)
img_new = img_new.permute(0, 3, 2, 1)
out = self.model(img_new)
return out.argmax(), img
class ObjectDetection(ModelComponent):
def __init__(self, model):
self.model = model
@expose(
inputs={"img": Image()},
outputs={
"boxes": Repeated(BBox()),
"labels": Repeated(Label("classes.txt"))
},
)
def detect(self, img):
img = img.permute(0, 3, 2, 1).float() / 255
out = self.model(img)[0]
return out["boxes"], out["labels"]
class ClassificationInferenceModelSequence(ModelComponent):
def __init__(self, model):
self.model1 = model[0]
self.model2 = model[1]
@expose(
inputs={"img": Image(extension="JPG")},
outputs={"prediction": Label(path=str(CWD / "imagenet_labels.txt"))},
)
def classify(self, img):
img = img.float() / 255
mean = torch.tensor([[[0.485, 0.456, 0.406]]]).float()
std = torch.tensor([[[0.229, 0.224, 0.225]]]).float()
img = (img - mean) / std
img = img.permute(0, 3, 2, 1)
out = self.model1(img)
out2 = self.model2(img)
assert out.argmax() == out2.argmax()
return out.argmax()
class ClassificationInferenceRepeated(ModelComponent):
def __init__(self, model):
self.model = model
@expose(
inputs={"img": Repeated(Image(extension="JPG"))},
outputs={
"prediction": Repeated(Label(path=str(CWD / "imagenet_labels.txt"))),
"other": Number(),
},
)
def classify(self, img):
img = img[0].float() / 255
mean = torch.tensor([[[0.485, 0.456, 0.406]]]).float()
std = torch.tensor([[[0.229, 0.224, 0.225]]]).float()
img = (img - mean) / std
img = img.permute(0, 3, 2, 1)
out = self.model(img)
return ([out.argmax(), out.argmax()], torch.Tensor([21]))
def test_path(session_global_datadir):
label = Label(path=str(session_global_datadir / "imagenet_labels.txt"))
assert label.deserialize("chickadee") == torch.tensor(19)
assert label.serialize(torch.tensor(19)) == "chickadee"
def test_wrong_type():
with pytest.raises(TypeError):
Label(classes=set())
with pytest.raises(ValueError):
Label(classes=None)
def test_dict():
label = Label(classes={56: "classA", 48: "classB"})
assert label.deserialize("classA") == torch.tensor(56)
with pytest.raises(TypeError):
Label(classes={"wrongtype": "classA"})
def test_list():
label = Label(classes=["classA", "classB"])
assert label.deserialize("classA") == torch.tensor(0)