keras-io
/

semantic-segmentation

Image Segmentation

Model card Files Files and versions Community

merve HF staff commited on Nov 26, 2021

Commit

7aad423

·

1 Parent(s): 86c02a6

Update pipeline.py

Files changed (1) hide show

pipeline.py +3 -3

pipeline.py CHANGED Viewed

@@ -57,11 +57,11 @@ class PreTrainedPipeline():
             mask_codes[f"mask_{cls}"] = mask
-        for i in range(pred_mask.shape[-1]): #for every class
             labels.append({
-                "label": f"LABEL_{i}",
-                "mask": mask_codes[f"mask_{i}"],
                 "score": 1.0,
             })
         return labels

             mask_codes[f"mask_{cls}"] = mask
             labels.append({
+                "label": f"LABEL_{cls}",
+                "mask": mask_codes[f"mask_{cls}"],
                 "score": 1.0,
             })
         return labels