fix #630: Add input size check to model_inference (#633)

* fix #630: Add input size check to model_inference

* fix #630: Added exception when receiving empty image
This commit is contained in:
mpena-vina 2021-12-02 04:00:23 +01:00 committed by GitHub
parent 69d5040590
commit d3d56182c6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 0 deletions

View File

@ -98,6 +98,8 @@ def model_inference(model,
if isinstance(imgs, (list, tuple)):
is_batch = True
if len(imgs) == 0:
raise Exception('empty imgs provided, please check and try again')
if not isinstance(imgs[0], (np.ndarray, str)):
raise AssertionError('imgs must be strings or numpy arrays')

View File

@ -102,3 +102,20 @@ def test_model_batch_inference_recog(cfg_file):
results = model_inference(model, [img, img], batch_mode=True)
assert len(results) == 2
@pytest.mark.parametrize(
'cfg_file',
['../configs/textdet/psenet/psenet_r50_fpnf_600e_icdar2017.py'])
def test_model_batch_inference_empty_detection(cfg_file):
tmp_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
config_file = os.path.join(tmp_dir, cfg_file)
model = build_model(config_file)
empty_detection = []
with pytest.raises(
Exception,
match='empty imgs provided, please check and try again'):
model_inference(model, empty_detection, batch_mode=True)