Major re-write to significantly improve UI

app.py

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

inferencing.py

@@ -9,36 +9,39 @@ import dotenv
 dotenv.load_dotenv()
 hf_tk = os.getenv('HF_AT')
 
-# Function to display detections and segmentation masks
 def display_detections(frame, results):
-    detections = []
+    class_names = []
+    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)]
-                detections.append([class_name, conf, x1, y1, x2, y2])
+                if class_name not in class_names:
+                    class_names.append(class_name)
+                annotations.append(((int(x1), int(y1), int(x2), int(y2)), class_name))
             outframe = results[0].plot()
-            return outframe, detections
+            return outframe, annotations, class_names
 
     print("No detections")
-    return frame, detections
-
+    return frame, [], []
 
 def process_image(input_image):
     results = model(input_image)
-    processed_frame, detections = display_detections(input_image, results)
+    return display_detections(input_image, results)
 
-    if not detections:  # Check if detections list is empty
+def get_habitat_data(class_names):
+    if not class_names:
         inconclusive_res = get_inconclusive()
-        return processed_frame, pd.DataFrame({'Message': ['No classes detected']}),  inconclusive_res[0], inconclusive_res[1]
+        return  inconclusive_res[0], inconclusive_res[1]
     else:
-        df = pd.DataFrame(detections, columns=["Class", "Confidence", "X1", "Y1", "X2", "Y2"])
-        habitat_res = classify_habitat(df['Class'].tolist())
-        return processed_frame, df[['Class', 'Confidence']], habitat_res[0], habitat_res[1]
-
-
+        habitat_res = classify_habitat(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
 # 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)