Merge branch 'main' of https://huggingface.co/merve/pet-segmentation

README.md

@@ -1,6 +1,7 @@
 ---
 tags:
 - image-segmentation
+- generic
 dataset:
 - oxfort-iit pets
 license: apache-2.0

pipeline.py

@@ -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}"],
+                "label": f"LABEL_{cls}",
+                "mask": mask_codes[f"mask_{cls}"],
                 "score": 1.0,
             })
         return labels