Update app.py

app.py

@@ -9,11 +9,14 @@ import torchvision.transforms.functional as F
 from ultralytics import YOLO
 import gradio as gr
 from huggingface_hub import hf_hub_download
+from datasets import load_dataset
 
 model_path_1 = hf_hub_download("Axpo/yolov8x-without-insulators", "yolov8x-without-insulators.pt", use_auth_token=os.environ["TOKEN"])
 model_path_2 = hf_hub_download("Axpo/yolov8x-with-insulators", "yolov8x-with-insulators.pt", use_auth_token=os.environ["TOKEN"])
 model_paths = [model_path_1, model_path_2]
 
+dataset = load_dataset("Axpo/example_pole_images", use_auth_token=os.environ["TOKEN"])
+
 def draw_bboxes_yolo(
         image: torch.Tensor, 
         bboxes: torch.Tensor, 
@@ -112,14 +115,7 @@ with gr.Blocks() as demo:
             )
 
     gr.Examples(
-        [
-            Path(os.path.dirname(__file__)) / "images" / "steel-pole-close.jpg",
-            Path(os.path.dirname(__file__)) / "images" / "steel-pole-overview.jpeg",
-            Path(os.path.dirname(__file__)) / "images" / "steel-pole-topdown.jpeg",
-            Path(os.path.dirname(__file__)) / "images" / "steel-pole-insulators.jpeg",
-            Path(os.path.dirname(__file__)) / "images" / "concrete-pole-close.jpeg",
-            Path(os.path.dirname(__file__)) / "images" / "concrete-pole-close2.jpeg",
-            ],
+        dataset["train"]["image"],
         input_img,
         output_img,
         predict_big_components