Merge pull request #6 from guigarfr/inference_multi-dataset

Adapt inference to multiple datasets
guigarfr · Oct 27, 2021 · 333d9cc · 333d9cc
2 parents 6c55265 + f658d8b
commit 333d9cc
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/mmdet/apis/inference.py b/mmdet/apis/inference.py
@@ -108,14 +108,17 @@ def inference_detector(model, imgs):
 
     cfg = model.cfg
     device = next(model.parameters()).device  # model device
+
+    if isinstance(cfg.data.test, list):
+        test_pipeline = cfg.data.test[0].pipeline
+    else:
+        test_pipeline = cfg.data.test.pipeline
 
     if isinstance(imgs[0], np.ndarray):
-        cfg = cfg.copy()
         # set loading pipeline type
-        cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam'
+        test_pipeline[0].type = 'LoadImageFromWebcam'
 
-    cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
-    test_pipeline = Compose(cfg.data.test.pipeline)
+    test_pipeline = pipelines.Compose(replace_ImageToTensor(test_pipeline))
 
     datas = []
     for img in imgs: