Tweak image labels

notebooks/app.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 179,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -43,7 +43,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 180,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 181,
+   "execution_count": 35,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -99,15 +99,17 @@
     "    grd = widgets.VBox([widgets.Label(label), img_widget])\n",
     "    return grd\n",
     "\n",
-    "def label_img(img, model) -> Image:\n",
-    "    pred = model.predict(img, device = 'cpu')[0].plot(labels = False)\n",
-    "    pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)\n",
-    "    return Image.fromarray(pred)"
+    "def label_img(img, model) -> (Image, int):\n",
+    "    pred = model.predict(img, device = 'cpu')[0]\n",
+    "    n = pred.boxes.data.shape[0]\n",
+    "    plot = pred.plot(labels = False)\n",
+    "    plot = cv2.cvtColor(plot, cv2.COLOR_BGR2RGB)\n",
+    "    return Image.fromarray(plot), n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 182,
+   "execution_count": 38,
    "metadata": {},
    "outputs": [
     {
@@ -115,11 +117,11 @@
      "output_type": "stream",
      "text": [
       "\n",
-      "0: 640x640 1 surveillance, 262.8ms\n",
-      "Speed: 1.9ms preprocess, 262.8ms inference, 1.1ms postprocess per image at shape (1, 3, 640, 640)\n",
+      "0: 640x640 2 surveillances, 430.8ms\n",
+      "Speed: 19.0ms preprocess, 430.8ms inference, 1.3ms postprocess per image at shape (1, 3, 640, 640)\n",
       "\n",
-      "0: 640x640 3 surveillances, 988.3ms\n",
-      "Speed: 1.8ms preprocess, 988.3ms inference, 0.6ms postprocess per image at shape (1, 3, 640, 640)\n"
+      "0: 640x640 2 surveillances, 1938.6ms\n",
+      "Speed: 2.5ms preprocess, 1938.6ms inference, 1.1ms postprocess per image at shape (1, 3, 640, 640)\n"
      ]
     }
    ],
@@ -166,8 +168,13 @@
     "grid[0, :] = widgets.VBox([location, size, heading, pitch, fov, button],\n",
     "                          layout = widgets.Layout(height = 'auto'))\n",
     "\n",
+    "def singular(x):\n",
+    "    return '' if x == 1 else 's'\n",
     "\n",
     "def button_click(b):\n",
+    "    for i in range(3):\n",
+    "        grid[1:, i] = widgets.Label('Loading...')\n",
+    "        \n",
     "    img = get_sv_img(location=location.value,\n",
     "                     size=size.value,\n",
     "                     heading=heading.value,\n",
@@ -175,10 +182,14 @@
     "                     fov=fov.value)\n",
     "    if img is not None:\n",
     "        grid[1:, 0] = make_img_widget(img, 'Original image')\n",
-    "        yolo_pred = label_img(img, yolo)\n",
-    "        grid[1:, 1] = make_img_widget(yolo_pred, 'YOLO predictions')\n",
-    "        detr_pred = label_img(img, detr)\n",
-    "        grid[1:, 2] = make_img_widget(detr_pred, 'RT-DETR predictions')\n",
+    "        \n",
+    "        yolo_pred, yolo_n = label_img(img, yolo)\n",
+    "        yolo_suffix = singular(yolo_n)\n",
+    "        grid[1:, 1] = make_img_widget(yolo_pred, f'YOLO predictions ({(yolo_n)} result{yolo_suffix})')\n",
+    "        \n",
+    "        detr_pred, detr_n = label_img(img, detr)\n",
+    "        detr_suffix = singular(detr_n)\n",
+    "        grid[1:, 2] = make_img_widget(detr_pred, f'RT-DETR predictions ({detr_n} result{detr_suffix})')\n",
     "        \n",
     "\n",
     "button.on_click(button_click)\n"
@@ -186,13 +197,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 184,
+   "execution_count": 37,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9020055a57294d38b6f3e40659b12a1d",
+       "model_id": "f3343b2adf0f49dea6c1830d0cdbc9bb",
        "version_major": 2,
        "version_minor": 0
       },
@@ -207,6 +218,11 @@
    "source": [
     "display(grid)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {