Show multiple test images

roboflow · Oct 4, 2024 · 59f47a1 · 59f47a1
1 parent cb132e4
commit 59f47a1
Showing 1 changed file with 24 additions and 22 deletions.
diff --git a/notebooks/train-yolo11-object-detection-on-custom-dataset.ipynb b/notebooks/train-yolo11-object-detection-on-custom-dataset.ipynb
@@ -923,8 +923,6 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# Run inference on your model on a persistant, auto-scaling, cloud API\n",
-        "\n",
         "import os, random, cv2\n",
         "import supervision as sv\n",
         "import IPython\n",
@@ -933,27 +931,31 @@
         "model_id = project.id.split(\"/\")[1] + \"/\" + dataset.version\n",
         "model = inference.get_model(model_id, userdata.get('ROBOFLOW_API_KEY'))\n",
         "\n",
-        "# Choose random test set image\n",
+        "# Location of test set images\n",
         "test_set_loc = dataset.location + \"/test/images/\"\n",
-        "random_test_image = random.choice(os.listdir(test_set_loc))\n",
-        "print(\"running inference on \" + random_test_image)\n",
-        "\n",
-        "image = cv2.imread(test_set_loc + random_test_image)\n",
-        "\n",
-        "results = model.infer(image, confidence=0.4, overlap=30)[0]\n",
-        "detections = sv.Detections.from_inference(results)\n",
-        "\n",
-        "box_annotator = sv.BoxAnnotator()\n",
-        "label_annotator = sv.LabelAnnotator()\n",
-        "\n",
-        "annotated_image = box_annotator.annotate(\n",
-        "    scene=image, detections=detections)\n",
-        "annotated_image = label_annotator.annotate(\n",
-        "    scene=annotated_image, detections=detections)\n",
-        "\n",
-        "_,ret = cv2.imencode('.jpg', annotated_image) \n",
-        "i = IPython.display.Image(data=ret)\n",
-        "IPython.display.display(i)"
+        "test_images = os.listdir(test_set_loc)\n",
+        "\n",
+        "# Run inference on 4 random test images, or fewer if fewer images are available\n",
+        "for img_name in random.sample(test_images, min(4, len(test_images))):\n",
+        "    print(\"Running inference on \" + img_name)\n",
+        "    \n",
+        "    # Load image\n",
+        "    image = cv2.imread(os.path.join(test_set_loc, img_name))\n",
+        "    \n",
+        "    # Perform inference\n",
+        "    results = model.infer(image, confidence=0.4, overlap=30)[0]\n",
+        "    detections = sv.Detections.from_inference(results)\n",
+        "    \n",
+        "    # Annotate boxes and labels\n",
+        "    box_annotator = sv.BoxAnnotator()\n",
+        "    label_annotator = sv.LabelAnnotator()\n",
+        "    annotated_image = box_annotator.annotate(scene=image, detections=detections)\n",
+        "    annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)\n",
+        "    \n",
+        "    # Display annotated image\n",
+        "    _, ret = cv2.imencode('.jpg', annotated_image)\n",
+        "    i = IPython.display.Image(data=ret)\n",
+        "    IPython.display.display(i)\n"
       ]
     },
     {