Made improvements to display confidence value in instance segmentation classes

app.py

@@ -5,7 +5,10 @@ from inferencing import run_inference
 iface = gr.Interface(
     fn=run_inference,
     inputs="image",
-    outputs=[gradio.Annotatedimage(label="Annotated Image"),gradio.DataFrame(label="Predicted Benthic Habitat",headers=["Prediction", "Confidence"]),gradio.DataFrame(label="Habitat Probabilities",headers=["Classification", "Confidence"])],
+    outputs=[gradio.Annotatedimage(label="Annotated Image"),
+             gradio.DataFrame(label="Predicted Benthic Habitat",headers=["Prediction", "Confidence"]),
+             gradio.DataFrame(label="Habitat Confidence Values",headers=["Classification", "Confidence"]),
+             gradio.DataFrame(label="Identified Classes Confidence Values",headers=["Classification", "Confidence"])],
     title="AMBULANT Benthic Identification System"
 )
 iface.launch()

habitat_classification.py

@@ -82,5 +82,5 @@ def get_inconclusive():
         'Confidence': 'N/a'})
     probabilities_df = pd.DataFrame({
         'Classification': ['Inconclusive'],
-        'Probability': 'N/a'})
+        'Confidence': 'N/a'})
     return result_df, probabilities_df

inferencing.py

@@ -10,18 +10,17 @@ dotenv.load_dotenv()
 hf_tk = os.getenv('HF_AT')
 
 def display_detections(frame, results):
-    class_names = []
+    detections = []
     annotations = []
     for result in results:
         if result.boxes.data.nelement() != 0:
             for detection in result.boxes.data:
                 x1, y1, x2, y2, conf, cls = detection.cpu().numpy()
                 class_name = result.names[int(cls)]
-                if class_name not in class_names:
-                    class_names.append(class_name)
+                detections.append([class_name, conf])
                 annotations.append(((int(x1), int(y1), int(x2), int(y2)), class_name))
             outframe = results[0].plot()
-            return outframe, annotations, class_names
+            return outframe, annotations, pd.DataFrame(detections, columns=["Class", "Confidence"])
 
     print("No detections")
     return frame, [], []
@@ -39,9 +38,13 @@ def get_habitat_data(class_names):
         return habitat_res[0], habitat_res[1]
 
 def run_inference(input_image):
-    processed_frame, annotations, class_names = process_image(input_image)
-    habitat, probabilities = get_habitat_data(class_names)
-    return (processed_frame, annotations), habitat, probabilities
+    processed_frame, annotations, detections = process_image(input_image)
+    if detections.empty:
+        habitat, probabilities = get_inconclusive()
+        detections = pd.DataFrame(columns=["Class", "Confidence"])
+    else:
+        habitat, probabilities = get_habitat_data(detections['Class'].tolist())
+    return (processed_frame, annotations), habitat, probabilities, detections
 # Load YOLO model
 raw_model = hf_hub_download(repo_id="dev-deg/Ambulant_1.0", filename="latest.pt",use_auth_token=hf_tk)
 model = YOLO(raw_model)