Simplify notebook further.

Author: davidbau

notebooks/DemoSegmenter.ipynb

@@ -78,9 +78,13 @@
     "    for row in reader:\n",
     "        names[int(row[0])] = row[5].split(\";\")[0]\n",
     "\n",
-    "def visualize_result(data, pred):\n",
-    "    (img, info) = data\n",
-    "\n",
+    "def visualize_result(img, pred, index=None):\n",
+    "    # filter prediction class if requested\n",
+    "    if index is not None:\n",
+    "        pred = pred.copy()\n",
+    "        pred[pred != index] = -1\n",
+    "        print(f'{names[index+1]}:')\n",
+    "        \n",
     "    # colorize prediction\n",
     "    pred_color = colorEncode(pred, colors).astype(numpy.uint8)\n",
     "\n",
@@ -146,8 +150,9 @@
     "        mean=[0.485, 0.456, 0.406], # These are RGB mean+std values\n",
     "        std=[0.229, 0.224, 0.225])  # across a large photo dataset.\n",
     "])\n",
-    "img_data = pil_to_tensor(\n",
-    "    Image.open('ADE_val_00001519.jpg').convert('RGB'))\n",
+    "pil_image = Image.open('ADE_val_00001519.jpg').convert('RGB')\n",
+    "img_original = numpy.array(pil_image)\n",
+    "img_data = pil_to_tensor(pil_image)\n",
     "singleton_batch = {'img_data': img_data[None].cuda()}\n",
     "output_size = img_data.shape[1:]"
    ]
@@ -168,7 +173,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": false
+   },
    "outputs": [],
    "source": [
     "# Run the segmentation at the highest resolution.\n",
@@ -177,20 +184,17 @@
     "    \n",
     "# Get the predicted scores for each pixel\n",
     "_, pred = torch.max(scores, dim=1)\n",
-    "visualize_result(\n",
-    "      (dataset_test[0]['img_ori'], dataset_test[0]['info']),\n",
-    "      pred.cpu()[0].numpy())"
+    "pred = pred.cpu()[0].numpy()\n",
+    "visualize_result(img_original, pred)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Run the model at multiple sizes\n",
+    "## Showing classes individually\n",
     "\n",
-    "One way to get slightly cleaner predictions from a segmentation model is to run the model several times on the same image at different resolutions, and then take the average of the scores for prredictions.\n",
-    "\n",
-    "This code does that."
+    "To see which colors are which, here we visualize individual classes, one at a time."
    ]
   },
   {
@@ -199,55 +203,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# The following code averages segmenter scores at multiple resolutions for better results\n",
-    "def test(segmentation_module, loader, gpu):\n",
-    "    segmentation_module.eval()\n",
-    "\n",
-    "    for batch_data in loader:\n",
-    "        # process data\n",
-    "        batch_data = batch_data[0]\n",
-    "        segSize = (batch_data['img_ori'].shape[0],\n",
-    "                   batch_data['img_ori'].shape[1])\n",
-    "        img_resized_list = batch_data['img_data']\n",
-    "\n",
-    "        with torch.no_grad():\n",
-    "            scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1])\n",
-    "            scores = async_copy_to(scores, gpu)\n",
-    "\n",
-    "            for img in img_resized_list:\n",
-    "                feed_dict = batch_data.copy()\n",
-    "                feed_dict['img_data'] = img\n",
-    "                del feed_dict['img_ori']\n",
-    "                del feed_dict['info']\n",
-    "                feed_dict = async_copy_to(feed_dict, gpu)\n",
-    "\n",
-    "                # forward pass\n",
-    "                pred_tmp = segmentation_module(feed_dict, segSize=segSize)\n",
-    "                scores = scores + pred_tmp / len(cfg.DATASET.imgSizes)\n",
-    "\n",
-    "            _, pred = torch.max(scores, dim=1)\n",
-    "            pred = as_numpy(pred.squeeze(0).cpu())\n",
-    "\n",
-    "        # visualization\n",
-    "        visualize_result(\n",
-    "            (batch_data['img_ori'], batch_data['info']),\n",
-    "            pred\n",
-    "        )\n",
-    "        \n",
-    "gpu = 0\n",
-    "torch.cuda.set_device(gpu)\n",
-    "\n",
-    "loader_test = torch.utils.data.DataLoader(\n",
-    "    dataset_test,\n",
-    "    batch_size=1,\n",
-    "    shuffle=False,\n",
-    "    collate_fn=user_scattered_collate,\n",
-    "    num_workers=5,\n",
-    "    drop_last=False)\n",
-    "\n",
-    "segmentation_module.cuda()\n",
-    "\n",
-    "test(segmentation_module, loader_test, gpu)\n"
+    "# Top classes in answer\n",
+    "predicted_classes = numpy.bincount(pred.flatten()).argsort()[::-1]\n",
+    "for c in predicted_classes[:15]:\n",
+    "    visualize_result(img_original, pred, c)"
    ]
   }
  ],