evaluate_deep_fruit_vision.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Deep Fruit Vision Evaluation\n",
    "\n",
    "Our detection, ripeness, and defect models are all trained and tested on different datasets. This notebook evaluates the overall performance of Deep Fruit Vision on a unique, hand-labelled test dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "import torch\n",
    "from torch.utils.data import Subset\n",
    "from deepfruitvision import DeepFruitVision\n",
    "from modules.datasets import EnsembleDataset, save_dataset\n",
    "from torchvision.ops import box_iou\n",
    "\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "best_classification_weights = os.path.join('weights', 'detection', 'best_classification.pt')\n",
    "best_detection_weights = os.path.join('weights', 'detection', 'best_detection.pt')\n",
    "ripeness_weights = os.path.join('weights', 'ripeness', 'ripeness_model_fine_tuned')\n",
    "defect_weights = os.path.join('weights', 'defect', 'defect_model_fine_tuned')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can use the EnsembleDataset class to easily load images and labels from the test dataset. We reserve a bit of the ensemble dataset for fine-tuning our models. We use the seed to make sure that we get the same images across runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "eval_dir = os.path.join('dataset', 'deepfruitvision_eval')\n",
    "detection_eval_img_dir = os.path.join(eval_dir, 'images')\n",
    "detection_eval_label_dir = os.path.join(eval_dir, 'labels')\n",
    "\n",
    "os.makedirs(detection_eval_img_dir, exist_ok=True)\n",
    "os.makedirs(detection_eval_label_dir, exist_ok=True)\n",
    "\n",
    "seed = 123\n",
    "np.random.seed(seed)\n",
    "num_fine_tune_samples = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ensemble_dataset = EnsembleDataset('dataset', for_yolov5=True)\n",
    "\n",
    "random_indices = np.random.permutation(len(ensemble_dataset))\n",
    "fine_tune_indices = random_indices[:num_fine_tune_samples]\n",
    "test_indices = random_indices[num_fine_tune_samples:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving dataset to dataset\\deepfruitvision_eval\\images and dataset\\deepfruitvision_eval\\labels: 100%|██████████| 105/105 [00:00<00:00, 145.52it/s]\n"
     ]
    }
   ],
   "source": [
    "yolov5_test_dataset = Subset(ensemble_dataset, test_indices)\n",
    "save_dataset(yolov5_test_dataset, detection_eval_img_dir, detection_eval_label_dir) # save the test split of the ensemble dataset with Yolo-v5 labels to the disk\n",
    "\n",
    "ensemble_dataset.for_yolov5 = False\n",
    "deepfruitvision_test_dataset = Subset(ensemble_dataset, test_indices) # then just get the test split of the ensemble dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Detection mAP\n",
    "\n",
    "The detection/Yolo-v5 model is responsible for detecting and classifying fruits, so we can just use Yolo-v5's built-in evaluation script. We've provided two trained Yolo-v5 models. The first model, `best_detection_weights.pt` was trained on the detection dataset. The second model, `best_classification_weights.pt` was trained on the detection dataset and fine-tuned on images from the ensemble dataset. The `best_detection_weights.pt` model is good at detecting fruits, but it can have issues classifying them if the fruit is in a different environment than one it saw in training (such as the ensemble dataset). It is a good model to use as a base for fine-tuning.\n",
    "\n",
    "The following two tests are each Yolo-v5 model on the detection and classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mval: \u001B[0mdata=fine_tune_apple_papaya_mango.yaml, weights=['weights\\\\detection\\\\best_detection.pt'], batch_size=32, imgsz=416, conf_thres=0.001, iou_thres=0.6, max_det=300, task=test, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=yolov5\\runs\\val, name=exp, exist_ok=False, half=False, dnn=False\n",
      "YOLOv5  2022-11-21 Python-3.10.5 torch-1.12.0 CUDA:0 (NVIDIA GeForce GTX 1070, 8192MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs\n",
      "\u001B[34m\u001B[1mtest: \u001B[0mScanning C:\\Users\\jorda\\Documents\\School\\CMPE 295\\fruit-detection\\dataset\\deepfruitvision_eval\\labels... 105 images, 0 backgrounds, 0 corrupt: 100%|██████████| 105/105 00:05\n",
      "\u001B[34m\u001B[1mtest: \u001B[0mNew cache created: C:\\Users\\jorda\\Documents\\School\\CMPE 295\\fruit-detection\\dataset\\deepfruitvision_eval\\labels.cache\n",
      "                 Class     Images  Instances          P          R      mAP50   mAP50-95: 100%|██████████| 4/4 00:03\n",
      "                   all        105        417      0.234      0.513      0.278      0.193\n",
      "                 apple        105        202       0.26      0.371      0.326      0.235\n",
      "                papaya        105         81     0.0794      0.914        0.2       0.13\n",
      "                 mango        105        134      0.362      0.254      0.307      0.214\n",
      "Speed: 0.3ms pre-process, 5.9ms inference, 9.6ms NMS per image at shape (32, 3, 416, 416)\n",
      "Results saved to \u001B[1myolov5\\runs\\val\\exp40\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "%run yolov5/val.py --data fine_tune_apple_papaya_mango.yaml --weights {best_detection_weights} --img 416 --task test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mval: \u001B[0mdata=fine_tune_apple_papaya_mango.yaml, weights=['weights\\\\detection\\\\best_classification.pt'], batch_size=32, imgsz=416, conf_thres=0.001, iou_thres=0.6, max_det=300, task=test, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=yolov5\\runs\\val, name=exp, exist_ok=False, half=False, dnn=False\n",
      "YOLOv5  2022-11-21 Python-3.10.5 torch-1.12.0 CUDA:0 (NVIDIA GeForce GTX 1070, 8192MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs\n",
      "\u001B[34m\u001B[1mtest: \u001B[0mScanning C:\\Users\\jorda\\Documents\\School\\CMPE 295\\fruit-detection\\dataset\\deepfruitvision_eval\\labels.cache... 105 images, 0 backgrounds, 0 corrupt: 100%|██████████| 105/105 00:00\n",
      "                 Class     Images  Instances          P          R      mAP50   mAP50-95: 100%|██████████| 4/4 00:03\n",
      "                   all        105        417       0.85      0.815        0.9      0.667\n",
      "                 apple        105        202      0.865      0.918      0.936      0.744\n",
      "                papaya        105         81      0.842      0.722      0.871      0.544\n",
      "                 mango        105        134      0.844      0.806      0.895      0.715\n",
      "Speed: 0.4ms pre-process, 6.0ms inference, 7.9ms NMS per image at shape (32, 3, 416, 416)\n",
      "Results saved to \u001B[1myolov5\\runs\\val\\exp46\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "%run yolov5/val.py --data fine_tune_apple_papaya_mango.yaml --weights {best_classification_weights} --img 416 --task test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The `best_classification_weights.pt` model is much better at classifying the fruit in the ensemble dataset. The `best_detection_weights.pt` model struggles on the ensemble dataset, but that's because it it is bad at classifying the fruit in the dataset since there is a large distribution shift from the detection dataset. However, when we only focus on detection, we see that the `best_detection_weights.pt` model has a competitive mAP and generalizes well to the ensemble dataset, despite not having seen any fruit from it. The `best_classification_weights.pt` model still does better, but that' due to the fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mval: \u001B[0mdata=fine_tune_apple_papaya_mango.yaml, weights=['weights\\\\detection\\\\best_detection.pt'], batch_size=32, imgsz=416, conf_thres=0.001, iou_thres=0.6, max_det=300, task=test, device=, workers=8, single_cls=True, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=yolov5\\runs\\val, name=exp, exist_ok=False, half=False, dnn=False\n",
      "YOLOv5  2022-11-21 Python-3.10.5 torch-1.12.0 CUDA:0 (NVIDIA GeForce GTX 1070, 8192MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs\n",
      "\u001B[34m\u001B[1mtest: \u001B[0mScanning C:\\Users\\jorda\\Documents\\School\\CMPE 295\\fruit-detection\\dataset\\deepfruitvision_eval\\labels.cache... 105 images, 0 backgrounds, 0 corrupt: 100%|██████████| 105/105 00:00\n",
      "                 Class     Images  Instances          P          R      mAP50   mAP50-95: 100%|██████████| 4/4 00:03\n",
      "                   all        105        417      0.787      0.835      0.878      0.651\n",
      "Speed: 0.3ms pre-process, 5.3ms inference, 8.5ms NMS per image at shape (32, 3, 416, 416)\n",
      "Results saved to \u001B[1myolov5\\runs\\val\\exp47\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "%run yolov5/val.py --data fine_tune_apple_papaya_mango.yaml --weights {best_detection_weights} --img 416 --task test --single-cls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mval: \u001B[0mdata=fine_tune_apple_papaya_mango.yaml, weights=['weights\\\\detection\\\\best_classification.pt'], batch_size=32, imgsz=416, conf_thres=0.001, iou_thres=0.6, max_det=300, task=test, device=, workers=8, single_cls=True, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=yolov5\\runs\\val, name=exp, exist_ok=False, half=False, dnn=False\n",
      "YOLOv5  2022-11-21 Python-3.10.5 torch-1.12.0 CUDA:0 (NVIDIA GeForce GTX 1070, 8192MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs\n",
      "\u001B[34m\u001B[1mtest: \u001B[0mScanning C:\\Users\\jorda\\Documents\\School\\CMPE 295\\fruit-detection\\dataset\\deepfruitvision_eval\\labels.cache... 105 images, 0 backgrounds, 0 corrupt: 100%|██████████| 105/105 00:00\n",
      "                 Class     Images  Instances          P          R      mAP50   mAP50-95: 100%|██████████| 4/4 00:02\n",
      "                   all        105        417       0.86        0.9      0.938       0.72\n",
      "Speed: 0.3ms pre-process, 5.2ms inference, 7.6ms NMS per image at shape (32, 3, 416, 416)\n",
      "Results saved to \u001B[1myolov5\\runs\\val\\exp49\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "%run yolov5/val.py --data fine_tune_apple_papaya_mango.yaml --weights {best_classification_weights} --img 416 --task test --single-cls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "And now we can clean the Yolov5 eval dir because we don't need it any more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# also, clean up the yolov5 eval dataset\n",
    "shutil.rmtree(eval_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Ensemble Classification Accuracy\n",
    "\n",
    "Now we want to evaluate the accuracy of the ensemble classification model. Since we only care about the accuracy of the harvestability, we ignore any bounding box that is too small to be classified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "YOLOv5  2022-11-21 Python-3.10.5 torch-1.12.0 CUDA:0 (NVIDIA GeForce GTX 1070, 8192MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs\n",
      "Adding AutoShape... \n"
     ]
    }
   ],
   "source": [
    "fruit_vision = DeepFruitVision(best_classification_weights, ripeness_weights, defect_weights, min_bounding_box_size=.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# use Matplotlib to display the true and predicted bounding boxes\n",
    "# true bounding boxes are a list of dicts with x, y, w, and h keys (all normalized, not pixel values)\n",
    "# predicted bounding boxes are a list of dicts with xmin, ymin, xmax, ymax (all normalized, not pixel values)\n",
    "def display_bboxes(img, true_bboxes, pred_bboxes):\n",
    "    img_h, img_w, _ = img.shape\n",
    "    fig, ax = plt.subplots(1)\n",
    "    ax.imshow(img)\n",
    "    for bbox in true_bboxes:\n",
    "        x = bbox['x'] * img_w\n",
    "        y = bbox['y'] * img_h\n",
    "        w = bbox['w'] * img_w\n",
    "        h = bbox['h'] * img_h\n",
    "        rect = matplotlib.patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='g', facecolor='none')\n",
    "        ax.add_patch(rect)\n",
    "    for bbox in pred_bboxes:\n",
    "        x = bbox['xmin'] * img_w\n",
    "        y = bbox['ymin'] * img_h\n",
    "        w = (bbox['xmax'] - bbox['xmin']) * img_w\n",
    "        h = (bbox['ymax'] - bbox['ymin']) * img_h\n",
    "        rect = matplotlib.patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none')\n",
    "        ax.add_patch(rect)\n",
    "    plt.show()\n",
    "\n",
    "for i in range(1):\n",
    "    img, true_boxes = deepfruitvision_test_dataset[i]\n",
    "    pred_boxes = fruit_vision.get_harvestability(img)\n",
    "    display_bboxes(img, true_boxes, pred_boxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "Evaluating DeepFruitVision:   0%|          | 0/105 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f6f93d8e02574856b1c74cbd82c87670"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 327 out of 417 correct (78.42%)\n",
      "Got 396 out of 417 correct bounding boxes (94.96%)\n",
      "Got 340 out of 417 correct ripeness labels (81.53%)\n",
      "Got 341 out of 417 correct ripeness labels when the bounding box was correct (81.77%)\n",
      "Got 364 out of 417 correct defect labels (87.29%)\n"
     ]
    }
   ],
   "source": [
    "min_box_size = fruit_vision.min_bounding_box_size\n",
    "\n",
    "total_boxes = 0\n",
    "total_correct = 0\n",
    "\n",
    "correct_bbox = 0\n",
    "correct_ripeness = 0\n",
    "correct_ripeness_true_bbox = 0\n",
    "correct_defect = 0\n",
    "\n",
    "for frame, true_boxes in tqdm(deepfruitvision_test_dataset, desc='Evaluating DeepFruitVision', total=len(deepfruitvision_test_dataset)):\n",
    "    predicted_boxes = fruit_vision.get_harvestability(frame)\n",
    "\n",
    "    # ignore any boxes in the predicted and true boxes that are too small to be classified\n",
    "    predicted_boxes = [box for box in predicted_boxes if box['xmax'] - box['xmin'] > min_box_size and box['ymax'] - box['ymin'] > min_box_size]\n",
    "    # the true boxes are numpy arrays with the format [class, x, y, w, h, ...]\n",
    "    true_boxes = [box for box in true_boxes if box['w'] > min_box_size and box['h'] > min_box_size]\n",
    "\n",
    "    if len(true_boxes) == 0: # if there are no boxes large enough to be classified, then we can't evaluate this frame\n",
    "        continue\n",
    "\n",
    "    total_boxes += len(true_boxes)\n",
    "\n",
    "    if len(predicted_boxes) == 0: # if there are no predicted boxes, then we automatically get 0 correct\n",
    "        continue\n",
    "    \n",
    "    # now we have to convert the predicted and true boxes to tensors so we can use Yolo-v5's built-in IoU function\n",
    "    predicted_boxes_tensor = torch.tensor([[box['xmin'], box['ymin'], box['xmax'], box['ymax']] for box in predicted_boxes])\n",
    "    # each box has to have the format [x1, y1, x2, y2] where x1 < x2 and y1 < y2\n",
    "    true_boxes_tensor = torch.tensor([[box['x'], box['y'], box['x'] + box['w'], box['y'] + box['h']] for box in true_boxes])\n",
    "\n",
    "    # this returns a tensor [num_true_boxes, num_predicted_boxes] that we can use to determine which of the true boxes have a corresponding predicted box\n",
    "    iou = box_iou(true_boxes_tensor, predicted_boxes_tensor)\n",
    "\n",
    "    # the max of each row will be the IoU of the predicted box with the true box\n",
    "    max_ious, max_iou_indices = torch.max(iou, dim=1)\n",
    "\n",
    "    for i, (max_iou, max_iou_index) in enumerate(zip(max_ious, max_iou_indices)):\n",
    "        if max_iou > 0.5: # if the IoU is greater than 0.5, then we consider it a correct prediction\n",
    "            correct_bbox += 1\n",
    "\n",
    "            if true_boxes[i]['ripeness'] == predicted_boxes[max_iou_index]['ripeness'][0]:\n",
    "                correct_ripeness += 1\n",
    "\n",
    "            if true_boxes[i]['defect'] == predicted_boxes[max_iou_index]['defect'][0]:\n",
    "                correct_defect += 1\n",
    "\n",
    "            true_ensemble_label = true_boxes[i]['ensemble']\n",
    "            predicted_harvestability_label = predicted_boxes[max_iou_index]['harvestability']\n",
    "\n",
    "            if true_ensemble_label == predicted_harvestability_label:\n",
    "                total_correct += 1\n",
    "\n",
    "print(f'Got {total_correct} out of {total_boxes} correct ({total_correct / total_boxes * 100:.2f}%)')\n",
    "print(f'Got {correct_bbox} out of {total_boxes} correct bounding boxes ({correct_bbox / total_boxes * 100:.2f}%)')\n",
    "print(f'Got {correct_ripeness} out of {total_boxes} correct ripeness labels ({correct_ripeness / total_boxes * 100:.2f}%)')\n",
    "print(f'Got {correct_defect} out of {total_boxes} correct defect labels ({correct_defect / total_boxes * 100:.2f}%)')"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Speed\n",
    "\n",
    "The time it takes to run the model on a single frame is variable, so we will run it on all the images in the ensemble dataset and take the average of them. We will load the images into memeory first so that we don't have to deal with the time to load them from the disk."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "data": {
      "text/plain": "Loading images:   0%|          | 0/155 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "be8fa11ffee14a1f87cad9a6ba80ddb6"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<matplotlib.image.AxesImage at 0x1a2378edea0>"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import time\n",
    "from modules.preprocessor import localize_fruit\n",
    "\n",
    "ensemble_dataset.for_yolov5 = False\n",
    "ensemble_images = []\n",
    "localized_img_tensors = []\n",
    "for frame, _ in tqdm(ensemble_dataset, desc='Loading images', total=len(ensemble_dataset)):\n",
    "    ensemble_images.append(frame)\n",
    "    bounding_boxes = fruit_vision.detection_module.get_bounding_boxes(frame)\n",
    "    localized_imgs = localize_fruit(frame, bounding_boxes, fruit_vision.min_bounding_box_size)\n",
    "\n",
    "    localized_img_tensors.append(localized_imgs)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "data": {
      "text/plain": "Running model:   0%|          | 0/155 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6d0ff58d06b742d6b7e15c7ced851d18"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "start = time.time()\n",
    "\n",
    "for frame in tqdm(ensemble_images, desc='Running model', total=len(ensemble_images)):\n",
    "    fruit_vision.get_harvestability(frame)\n",
    "\n",
    "end = time.time()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average time per frame: 0.69 seconds\n",
      "Average FPS: 1.46\n"
     ]
    }
   ],
   "source": [
    "print(f'Average time per frame: {(end - start) / len(ensemble_images):.2f} seconds')\n",
    "print(f'Average FPS: {len(ensemble_images) / (end - start):.2f}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We achieve around 1.5 FPS on each frame, or about .66 seconds. About half of that time is spent on the deep learning models (Yolo-v5, EfficientNet-v2S for defect and ripeness detection) and the other half is spent using the bounding boxes to localize the fruit for ripeness and defect detection. The `preprocessor.localize_fruit` is unoptimized and has a lot of room for improvement."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "outputs": [
    {
     "data": {
      "text/plain": "Running model:   0%|          | 0/155 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "cafd116b55414b1da6b7ef7a60f39641"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "start = time.time()\n",
    "\n",
    "for frame, tensor_stack in tqdm(zip(ensemble_images, localized_img_tensors), desc='Running model', total=len(localized_img_tensors)):\n",
    "    bounding_boxes = fruit_vision.detection_module.get_bounding_boxes(frame)\n",
    "    fruit_vision.ripeness_module.model(tensor_stack)\n",
    "    fruit_vision.defect_module.model(tensor_stack)\n",
    "\n",
    "end = time.time()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average time per frame: 0.32 seconds\n",
      "Average FPS: 3.13\n"
     ]
    }
   ],
   "source": [
    "print(f'Average time per frame: {(end - start) / len(ensemble_images):.2f} seconds')\n",
    "print(f'Average FPS: {len(ensemble_images) / (end - start):.2f}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}