Update app.py

app.py

@@ -15,15 +15,7 @@ model_path_1 = hf_hub_download("Axpo/yolov8x-without-insulators", "yolov8x-witho
 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]
 
-image_names = ["concrete-pole-close.jpeg",
-               "concrete-pole-close2.jpeg",
-               "steel-pole-close.jpg",
-               "steel-pole-insulators.jpeg",
-               "steel-pole-overview.jpeg",
-               "steel-pole-topdown.jpeg"]
-
-img_paths = [hf_hub_download("Axpo/example_pole_images", img_name, repo_type="dataset", use_auth_token=os.environ["TOKEN"]) 
-             for img_name in image_names]
+dataset = load_dataset("Axpo/example_pole_images", use_auth_token=os.environ["TOKEN"])
 
 def draw_bboxes_yolo(
         image: torch.Tensor, 
@@ -123,7 +115,7 @@ with gr.Blocks() as demo:
             )
 
     gr.Examples(
-        img_paths,
+        dataset["train"]["image"],
         input_img,
         output_img,
         predict_big_components