added confidence of segmentation from notebook detr

app.py

@@ -3,6 +3,13 @@ Using as reference:
 - https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512
 - https://huggingface.co/spaces/chansung/segformer-tf-transformers/blob/main/app.py
 - https://huggingface.co/facebook/detr-resnet-50-panoptic
+# https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
+
+https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/DETR/DETR_panoptic_segmentation_minimal_example_(with_DetrFeatureExtractor).ipynb
+
+Additions
+- add shown labels as strings
+- show only animal masks (ask an nlp model?)
 """
 
 from transformers import DetrFeatureExtractor, DetrForSegmentation
@@ -168,26 +175,31 @@ def ade_palette():
         [102, 255, 0],
         [92, 0, 255],
     ]
+                                           
 
-feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50-panoptic')
-model = DetrForSegmentation.from_pretrained('facebook/detr-resnet-50-panoptic')
-
-# gradio components
-input = gr.inputs.Image()
-output = gr.outputs.Image()
-
-def predict_animal_mask(im):
+def predict_animal_mask(im,
+                        flag_high_confidence):
     image = Image.fromarray(im) # im: numpy array 3d: 480, 640, 3: to PIL Image
     image = image.resize((200,200)) #  PIL image # could I upsample output instead? better?
 
-    inputs = feature_extractor(images=image, return_tensors="pt") #pt=Pytorch, tf=TensorFlow
-    outputs = model(**inputs)
-    logits = outputs.logits # torch.Size([1, 100, 251])
+    # encoding is a dict with pixel_values and pixel_mask
+    encoding = feature_extractor(images=image, return_tensors="pt") #pt=Pytorch, tf=TensorFlow
+    outputs = model(**encoding) # odict with keys: ['logits', 'pred_boxes', 'pred_masks', 'last_hidden_state', 'encoder_last_hidden_state']
+    logits = outputs.logits # torch.Size([1, 100, 251]); why 251?
     bboxes = outputs.pred_boxes
-    masks = outputs.pred_masks # torch.Size([1, 100, 200, 200])
+    masks = outputs.pred_masks # torch.Size([1, 100, 200, 200]); for every pixel, score in each of the 100 classes? there is a mask per class
+
+    # keep only the masks with high confidence?--------------------------------
+    # compute the prob per mask (i.e., class), excluding the "no-object" class (the last one)
+    if flag_high_confidence:
+        prob_per_query = outputs.logits.softmax(-1)[..., :-1].max(-1)[0] # why logits last dim 251?
+        # threshold the confidence
+        keep = prob_per_query > 0.85
+    else:
+        keep = torch.ones(outputs.logits.shape[0:2], dtype=torch.bool)
 
     # postprocess the mask (numpy arrays)
-    label_per_pixel = torch.argmax(masks.squeeze(),dim=0).detach().numpy()
+    label_per_pixel = torch.argmax(masks[keep].squeeze(),dim=0).detach().numpy() # from the masks per class, select the highest per pixel
     color_mask = np.zeros(image.size+(3,))
     for lbl, color in enumerate(ade_palette()):
         color_mask[label_per_pixel==lbl,:] = color
@@ -198,12 +210,23 @@ def predict_animal_mask(im):
 
     return pred_img
 
+#######################################
+# get models from hugging face
+feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50-panoptic')
+model = DetrForSegmentation.from_pretrained('facebook/detr-resnet-50-panoptic')
+
+# gradio components -inputs
+gr_image_input = gr.inputs.Image()
+gr_checkbox_high_confidence = gr.inputs.Checkbox(False,
+                                                 label='disply high confidence only?')
+# gradio outputs
+gr_image_output = gr.outputs.Image() 
 
 ####################################################
 # Create user interface and launch
 gr.Interface(predict_animal_mask, 
-    inputs = input, 
-    outputs = output,
+    inputs = [gr_image_input,gr_checkbox_high_confidence],
+    outputs = gr_image_output,
     title = 'Animals* segmentation in images',
     description = "An animal* segmentation image webapp using DETR (End-to-End Object Detection) model with ResNet-50 backbone").launch()