28 lines
1.3 KiB
Markdown
28 lines
1.3 KiB
Markdown
The two tiny ONNX models are deterministic classifiers used by the image upload API tests.
|
|
They accept the same `[1, 224, 224, 3]` float input as the production nsfw model and return
|
|
`[Drawings, Hentai, Neutral, Porn, Sexy]` scores. One always returns Neutral = 1 and the
|
|
other always returns Porn = 1. They exercise the real nsfw inference path without requiring
|
|
image fixtures.
|
|
|
|
## Deterministic model fixtures
|
|
|
|
To regenerate them with Python's `onnx` package:
|
|
|
|
```python
|
|
from pathlib import Path
|
|
from onnx import TensorProto, checker, helper
|
|
|
|
for name, scores in [("safe", [0., 0., 1., 0., 0.]), ("unsafe", [0., 0., 0., 1., 0.])]:
|
|
tensor = helper.make_tensor("probabilities", TensorProto.FLOAT, [1, 5], scores)
|
|
graph = helper.make_graph(
|
|
[helper.make_node("Identity", ["probabilities"], ["output"])],
|
|
"constant classifier for upload endpoint tests",
|
|
[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 224, 224, 3])],
|
|
[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 5])],
|
|
[tensor],
|
|
)
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)], ir_version=8)
|
|
checker.check_model(model)
|
|
Path(f"tests/fixtures/image-upload-{name}.onnx").write_bytes(model.SerializeToString())
|
|
```
|