فهرست منبع

Merge pull request #884 from will-jl944/develop_jf

fix seg prediction bug when input is single image
FlyingQianMM 4 سال پیش
والد
کامیت
699e493b1b
1فایلهای تغییر یافته به همراه9 افزوده شده و 6 حذف شده
  1. 9 6
      dygraph/paddlex/cv/models/segmenter.py

+ 9 - 6
dygraph/paddlex/cv/models/segmenter.py

@@ -478,12 +478,15 @@ class BaseSegmenter(BaseModel):
         label_map = label_map.numpy().astype('uint8')
         score_map = outputs['score_map']
         score_map = score_map.numpy().astype('float32')
-        prediction = [{
-            'label_map': l,
-            'score_map': s
-        } for l, s in zip(label_map, score_map)]
-        if isinstance(img_file, (str, np.ndarray)):
-            prediction = prediction[0]
+        if isinstance(img_file, list) and len(img_file) > 1:
+            prediction = [{
+                'label_map': l,
+                'score_map': s
+            } for l, s in zip(label_map, score_map)]
+        elif isinstance(img_file, list):
+            prediction = [{'label_map': label_map, 'score_map': score_map}]
+        else:
+            prediction = {'label_map': label_map, 'score_map': score_map}
         return prediction
 
     def _preprocess(self, images, transforms, model_type):