Browse Source

fix when rec res is None (#3228)

Tingquan Gao 9 months ago
parent
commit
81f9debe92

+ 2 - 2
paddlex/inference/pipelines/pp_shitu_v2/pipeline.py

@@ -106,9 +106,9 @@ class ShiTuV2Pipeline(BasePipeline):
         single_img_res = {"input_path": input_data, "input_img": raw_img, "boxes": []}
         for i, obj in enumerate(det_res["boxes"]):
             rec_scores = rec_res["score"][i]
-            if isinstance(rec_scores, np.ndarray):
-                rec_scores = rec_scores.tolist()
+            rec_scores = rec_scores.tolist() if rec_scores is not None else [None]
             labels = rec_res["label"][i]
+            labels = labels.tolist() if labels is not None else [None]
             single_img_res["boxes"].append(
                 {
                     "labels": labels,

+ 5 - 1
paddlex/inference/pipelines/pp_shitu_v2/result.py

@@ -79,7 +79,11 @@ def draw_box(img, boxes):
         )
 
         # draw label
-        text = "{} {:.2f}".format(dt["label"], score)
+        if score is not None:
+            text = "{} {:.2f}".format(dt["label"], score)
+        else:
+            text = "{}".format(dt["label"])
+
         if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
             tw, th = draw.textsize(text, font=font)
         else: