Added logging of images

yolo-inference-project/src/app.py

@@ -5,13 +5,34 @@ from utils import draw_bounding_boxes
 import time
 import subprocess
 import re
+from datetime import datetime
+from roboflow import Roboflow
 
+# Hugging Face Dataset configuration
+HF_DATASET_NAME = "vincentb25/flagged-images"  # Replace "username" with your HF username
+PRIVATE_DATASET = False
+
+# TODO get API KEY  
+rf = Roboflow(api_key=os.environ["ROBOFLOW_API_KEY"])
+workspaceId = 'ezbot'
+projectId = '2025'
+project = rf.workspace(workspaceId).project(projectId)
+
+
+def save_image_to_roboflow(image_path, model_name, conf_thresh):
+    """Save flagged images to roboflow."""
+
+    project.upload(image_path)
+
+
+# Load available models
 def get_model_names():
     model_dir = "models"  # Update this path to your models directory
     return [f for f in os.listdir(model_dir) if f.endswith('.pt')]  # Assuming models are in .pt format
 
 model_names = get_model_names()
 
+# Sample images for testing
 image_paths= [['examples/smaller_many_cans.jpg', 'yolo11m.pt', 0.5],
               ['examples/gazebo_all.jpg', 'yolo11m.pt', 0.5],
               ['examples/sequence_unity.jpg', 'yolo11m.pt', 0.5],
@@ -24,6 +45,7 @@ def inference(image, model_name, conf_thresh):
     if not image:
         raise gr.Error("No image provided. Please upload an image.")
     
+    # Perform inference
     model = load_model("models/" + model_name)
     results = perform_inference(model, image)
     output_image = draw_bounding_boxes(results, conf_thresh)
@@ -40,9 +62,13 @@ def inference(image, model_name, conf_thresh):
             speed_str += "\nRunning on '" + proc_name + "'"
             break
 
-    print(speed)
-    return output_image,speed_str
+    # Save flagged image metadata
+    save_image_to_roboflow(image, model_name, conf_thresh)
+    print(f"Image flagged and saved: {image}")
+
+    return output_image, speed_str
 
+# Define the Gradio app
 demo = gr.Interface(
     fn=inference,
     inputs=[
@@ -57,7 +83,7 @@ demo = gr.Interface(
     title="YOLO Model Inference",
     description="Select a YOLO model, upload an image, and set the confidence threshold to perform inference.",
     examples=image_paths,
-    # flagging_mode="auto"
+    flagging_mode="auto"  # Automatically flag all inputs
 )
 
-demo.launch()
+demo.launch()