Add fallback labels for segmentation

Author: kipcole9

lib/segmentation.ex

@@ -75,6 +75,152 @@ if ImageVision.ortex_configured?() do
     @detr_short_side 800
     @detr_no_object_class 250
 
+    # Canonical COCO panoptic id → label map (133 categories, IDs 1-200
+    # with gaps), drawn from the official panoptic_coco_categories.json:
+    # https://github.com/cocodataset/panopticapi/blob/master/panoptic_coco_categories.json
+    #
+    # Used as a fallback for IDs where the model's config.json carries
+    # a placeholder `"LABEL_n"` instead of the real class name. Both the
+    # original `facebook/detr-resnet-50-panoptic` config and the Xenova
+    # ONNX repackage drop the names for IDs 183-200 (the COCO panoptic
+    # "merged" stuff classes — sky-other-merged, mountain-merged, etc.),
+    # even though the model predicts them confidently.
+    @coco_panoptic_id2label %{
+      1 => "person",
+      2 => "bicycle",
+      3 => "car",
+      4 => "motorcycle",
+      5 => "airplane",
+      6 => "bus",
+      7 => "train",
+      8 => "truck",
+      9 => "boat",
+      10 => "traffic light",
+      11 => "fire hydrant",
+      13 => "stop sign",
+      14 => "parking meter",
+      15 => "bench",
+      16 => "bird",
+      17 => "cat",
+      18 => "dog",
+      19 => "horse",
+      20 => "sheep",
+      21 => "cow",
+      22 => "elephant",
+      23 => "bear",
+      24 => "zebra",
+      25 => "giraffe",
+      27 => "backpack",
+      28 => "umbrella",
+      31 => "handbag",
+      32 => "tie",
+      33 => "suitcase",
+      34 => "frisbee",
+      35 => "skis",
+      36 => "snowboard",
+      37 => "sports ball",
+      38 => "kite",
+      39 => "baseball bat",
+      40 => "baseball glove",
+      41 => "skateboard",
+      42 => "surfboard",
+      43 => "tennis racket",
+      44 => "bottle",
+      46 => "wine glass",
+      47 => "cup",
+      48 => "fork",
+      49 => "knife",
+      50 => "spoon",
+      51 => "bowl",
+      52 => "banana",
+      53 => "apple",
+      54 => "sandwich",
+      55 => "orange",
+      56 => "broccoli",
+      57 => "carrot",
+      58 => "hot dog",
+      59 => "pizza",
+      60 => "donut",
+      61 => "cake",
+      62 => "chair",
+      63 => "couch",
+      64 => "potted plant",
+      65 => "bed",
+      67 => "dining table",
+      70 => "toilet",
+      72 => "tv",
+      73 => "laptop",
+      74 => "mouse",
+      75 => "remote",
+      76 => "keyboard",
+      77 => "cell phone",
+      78 => "microwave",
+      79 => "oven",
+      80 => "toaster",
+      81 => "sink",
+      82 => "refrigerator",
+      84 => "book",
+      85 => "clock",
+      86 => "vase",
+      87 => "scissors",
+      88 => "teddy bear",
+      89 => "hair drier",
+      90 => "toothbrush",
+      92 => "banner",
+      93 => "blanket",
+      95 => "bridge",
+      100 => "cardboard",
+      107 => "counter",
+      109 => "curtain",
+      112 => "door-stuff",
+      118 => "floor-wood",
+      119 => "flower",
+      122 => "fruit",
+      125 => "gravel",
+      128 => "house",
+      130 => "light",
+      133 => "mirror-stuff",
+      138 => "net",
+      141 => "pillow",
+      144 => "platform",
+      145 => "playingfield",
+      147 => "railroad",
+      148 => "river",
+      149 => "road",
+      151 => "roof",
+      154 => "sand",
+      155 => "sea",
+      156 => "shelf",
+      159 => "snow",
+      161 => "stairs",
+      166 => "tent",
+      168 => "towel",
+      171 => "wall-brick",
+      175 => "wall-stone",
+      176 => "wall-tile",
+      177 => "wall-wood",
+      178 => "water-other",
+      180 => "window-blind",
+      181 => "window-other",
+      184 => "tree-merged",
+      185 => "fence-merged",
+      186 => "ceiling-merged",
+      187 => "sky-other-merged",
+      188 => "cabinet-merged",
+      189 => "table-merged",
+      190 => "floor-other-merged",
+      191 => "pavement-merged",
+      192 => "mountain-merged",
+      193 => "grass-merged",
+      194 => "dirt-merged",
+      195 => "paper-merged",
+      196 => "food-other-merged",
+      197 => "building-other-merged",
+      198 => "rock-merged",
+      199 => "wall-other-merged",
+      200 => "rug-merged"
+    }
+
     @default_min_score 0.5
 
     # ImageNet normalisation (shared by both models).
@@ -549,6 +695,17 @@ if ImageVision.ortex_configured?() do
       {{tensor, pixel_mask}, input_h, input_w}
     end
 
+    # Resolves a class index to a human-readable label, preferring the
+    # repo's id2label and falling back to the canonical COCO panoptic
+    # map when the repo carries a `LABEL_n` placeholder.
+    defp lookup_panoptic_label(id2label, class_idx) do
+      case Map.get(id2label, to_string(class_idx)) do
+        nil -> Map.get(@coco_panoptic_id2label, class_idx, "class_#{class_idx}")
+        "LABEL_" <> _ -> Map.get(@coco_panoptic_id2label, class_idx, "class_#{class_idx}")
+        name -> name
+      end
+    end
+
     # Loads id2label from config.json; cached in :persistent_term.
     defp load_detr_labels(repo) do
       key = {__MODULE__, :labels, repo}
@@ -597,7 +754,7 @@ if ImageVision.ortex_configured?() do
       |> Enum.with_index()
       |> Enum.flat_map(fn {{class_idx, score}, query_idx} ->
         if class_idx != @detr_no_object_class and score >= min_score do
-          label = Map.get(id2label, to_string(class_idx), "class_#{class_idx}")
+          label = lookup_panoptic_label(id2label, class_idx)
           mask_tensor = pred_masks[0][query_idx]
 
           mask =