object_detection/detectron2_training-kfold.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting a (Detectron2) Object Detection Model via K-Fold Cross-Validation to produce Out-of-sample Predictions for every Image in a Dataset \n",
    "\n",
    "This notebook demonstrates how to fit a [Detectron2](https://github.com/facebookresearch/detectron2/) model on object detection datasets via *cross-validation*, and produce *out-of-sample* predictions required to run cleanlab for detecting label errors in such data.  Out-of-sample predictions (from a copy of the model that was never trained on the image it is producing a prediction for) are less subject to overfitting and can more accurately detect label issues, but are harder to obtain for every image in an entire dataset unless we employ K-fold cross-validation as demonstrated here.\n",
    "\n",
    "Thus run this notebook if you wish to find every mislabeled image an a dataset. Once you have the out-of-sample predictions produced from this notebook, you can directly run the cleanlab [tutorial](https://docs.cleanlab.ai/stable/index.html) notebook on [Finding Label Errors in Object Detection Datasets](https://github.com/cleanlab/cleanlab/blob/master/docs/source/tutorials/object_detection.ipynb). Even though the cleanlab tutorial is focused on finding issues in only the validation subset of COCO, you can easily update the tutorial to run on the full dataset for which we obtain predictions via cross-validation in this notebook.\n",
    "\n",
    "In object detection data, each image is annotated with multiple bounding boxes.  Each bounding box surrounds a physical object within an image scene, and is annotated with a given class label. Using this labeled data, we train a model to predict the locations and classes of objects in an image. The trained model can subsequently be used to identify mislabeled images, which when corrected, allow you to train an even better model without changing your training code! \n",
    "\n",
    "Here we fit a state-of-the-art neural network trained starting from a pretrained [X-101](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md#imagenet-pretrained-models) network backbone. First let's import the required packages and download the data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cleanlab/examples/blob/master/object_detection/detectron2_training.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from detectron2.engine import DefaultTrainer\n",
    "from detectron2.config import get_cfg\n",
    "import pickle\n",
    "# import some common libraries\n",
    "from detectron2.data import build_detection_test_loader, build_detection_train_loader\n",
    "import numpy as np\n",
    "import os, json, cv2, random\n",
    "from detectron2.data import build_detection_test_loader\n",
    "from detectron2.evaluation import COCOEvaluator, inference_on_dataset\n",
    "# import some common detectron2 utilities\n",
    "from detectron2 import model_zoo\n",
    "from detectron2.engine import DefaultPredictor\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "from detectron2.data import MetadataCatalog, DatasetCatalog\n",
    "from detectron2.data.datasets import register_coco_instances\n",
    "import glob\n",
    "from sklearn.model_selection import KFold\n",
    "import json\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wget -nc \"http://images.cocodataset.org/annotations/annotations_trainval2017.zip\"  && unzip -q -o annotations_trainval2017.zip\n",
    "!wget -nc \"http://images.cocodataset.org/zips/val2017.zip\" && unzip -q -o val2017.zip\n",
    "!wget -nc \"http://images.cocodataset.org/zips/train2017.zip\" && unzip -q -o train2017.zip\n",
    "!wget -nc \"https://cleanlab-public.s3.amazonaws.com/ObjectDetectionBenchmarking/tutorial/TRAIN_COCO_ALL_labels.pkl\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In case an error occurs in the preceding cell, kindly download the [COCO dataset](https://cocodataset.org/#download).\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "If you wish to train on a custom dataset, be sure to review the COCO dataset guidelines for formatting your data, which can be found on their [website](https://cocodataset.org/#format-data)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split data into k-folds\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code performs K-Fold cross-validation on a dataset in the COCO format.\n",
    "* It splits the dataset into multiple folds and creates separate training and test data for each fold. The code loads a JSON file containing the dataset annotations, extracts the image IDs and maps them to file names. \n",
    "* It then defines a function to split the data based on the fold indices. The function creates separate training and test data dictionaries, preserving common metadata. It iterates over the images and annotations, assigning them to the appropriate data split based on the image ID. \n",
    "* Finally we save the training and test data as separate JSON files and print information about each fold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from sklearn.model_selection import KFold\n",
    "from collections import defaultdict\n",
    "NUM_FOLDS = 5  # Reduce this to get faster results but potentially less accurate predictions \n",
    "\n",
    "# Load the JSON file\n",
    "data = json.load(open(\"annotations/instances_train2017.json\", 'rb'))\n",
    "image_data = data['images']\n",
    "\n",
    "# Get unique image IDs and create a mapping of image ID to file name\n",
    "image_ids = [int(image['file_name'].split(\".\")[0]) for image in image_data]\n",
    "image_map = {image['id']: image['file_name'] for image in image_data}\n",
    "\n",
    "def split_data(train_indices, test_indices):\n",
    "    train_data = defaultdict(list)\n",
    "    test_data = defaultdict(list)\n",
    "    common_metadata = ['info', 'licenses', 'categories']\n",
    "\n",
    "    for cm in common_metadata:\n",
    "        train_data[cm] = data[cm]\n",
    "        test_data[cm] = data[cm]\n",
    "\n",
    "    train_image_ids = set([image_ids[i] for i in train_indices])\n",
    "    test_image_ids = set([image_ids[i] for i in test_indices])\n",
    "\n",
    "    for image in image_data:\n",
    "        image_id = int(image['file_name'].split('.')[0])\n",
    "        if image_id in train_image_ids:\n",
    "            train_data['images'].append(image)\n",
    "        else:\n",
    "            test_data['images'].append(image)\n",
    "\n",
    "    for annotation in data['annotations']:\n",
    "        image_id = annotation['image_id']\n",
    "        if image_id in train_image_ids:\n",
    "            train_data['annotations'].append(annotation)\n",
    "        else:\n",
    "            test_data['annotations'].append(annotation)\n",
    "\n",
    "    return train_data, test_data\n",
    "\n",
    "def print_data_info(data_dict, fold):\n",
    "    images_count = len(data_dict['images'])\n",
    "    annotations_count = len(data_dict['annotations'])\n",
    "    print(f\"Number of images: {images_count}, Number of annotations: {annotations_count}\")\n",
    "\n",
    "    \n",
    "def unregister_coco_instances(name):\n",
    "    if name in DatasetCatalog.list():\n",
    "        DatasetCatalog.remove(name)\n",
    "        MetadataCatalog.remove(name)\n",
    "\n",
    "# Generate K-Fold cross-validation\n",
    "kf = KFold(n_splits=NUM_FOLDS)\n",
    "pairs = []\n",
    "for fold, (train_indices, test_indices) in enumerate(kf.split(image_ids)):\n",
    "    train_data, test_data = split_data(train_indices, test_indices)\n",
    "    train_file = f\"train_coco_{fold}_fold.json\"\n",
    "    test_file = f\"test_coco_{fold}_fold.json\"\n",
    "    # Unregister instances with the same names only if they exist\n",
    "    unregister_coco_instances(train_file)\n",
    "    unregister_coco_instances(test_file)\n",
    "    # Register COCO instances for training and validation. \n",
    "    # Note: The 'train2017' folder is retained as the base path for images.\n",
    "    register_coco_instances(train_file, {}, train_file, \"train2017\")\n",
    "    register_coco_instances(test_file, {}, test_file, \"train2017\")\n",
    "    pairs.append([train_file,test_file])\n",
    "    with open(train_file, 'w') as train_file:\n",
    "        json.dump(train_data, train_file)\n",
    "    with open(test_file, 'w') as test_file:\n",
    "        json.dump(test_data, test_file)\n",
    "    print(f\"Data info for training data fold {fold}:\")\n",
    "    print_data_info(train_data, fold)\n",
    "    print(f\"Data info for test data fold {fold}:\")\n",
    "    print_data_info(test_data, fold)\n",
    "    \n",
    "TRAIN_PATH = os.path.join(os.getcwd(),\"train2017\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the configuration settings for training an object detection model using Detectron2. The model architecture used in this example is \"faster_rcnn_X_101_32x8d_FPN_3x\" from the COCO-Detection model zoo. The training data is specified by the \"my_dataset_train\" dataset and validation data is specified by the \"my_dataset_val\" dataset which refer to COCO2017 train and val containing only the subset of labels specified before.\n",
    "\n",
    "The number of worker threads is set to 2 and the batch size is set to 2.\n",
    "The learning rate and maximum number of iterations are also specified. The model is initialized from the COCO-Detection model zoo and the output directory for the trained model is created. Finally, the configuration is passed to the DefaultTrainer class for training the object detection model.\n",
    "\n",
    "<strong>Note:</strong> The choice of the number of iterations is informed by the incorporation of [early stopping.](https://en.wikipedia.org/wiki/Early_stopping#:~:text=In%20machine%20learning%2C%20early%20stopping,training%20data%20with%20each%20iteration.) This technique monitors the validation loss throughout training, saving the model upon improvement and halting training if no progress is observed within a defined patience period. Early stopping aims to identify an optimal model iteration, mitigating the risk of overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Early_stopping(DefaultTrainer):\n",
    "    def __init__(self, cfg, early_stop_patience=5, model_checkpoint_path=\"model_checkpoint.pth\"):\n",
    "        super().__init__(cfg)\n",
    "        self.early_stop_patience = early_stop_patience\n",
    "        self.model_checkpoint_path = model_checkpoint_path\n",
    "        self.best_validation_loss = float('inf')\n",
    "        self.current_patience = 0\n",
    "\n",
    "    def build_train_loader(self, cfg):\n",
    "        return build_detection_train_loader(cfg)\n",
    "  \n",
    "    def data_loader_mapper(self, batch):\n",
    "        return batch\n",
    "\n",
    "    def run_hooks(self):\n",
    "        val_loss = self.validation()\n",
    "        if val_loss < self.best_validation_loss:\n",
    "            self.best_validation_loss = val_loss\n",
    "            self.current_patience = 0\n",
    "            self.save_checkpoint()\n",
    "        else:\n",
    "            self.current_patience += 1\n",
    "            if self.current_patience >= self.early_stop_patience:\n",
    "                self._trainer.save_checkpoint()\n",
    "                self._trainer.has_finished = True\n",
    "\n",
    "    def validation(self):\n",
    "        # Define evaluator here\n",
    "        evaluator = COCOEvaluator(self.cfg.DATASETS.TEST[0], self.cfg, True, output_dir=\"./output/\")\n",
    "        val_loader = build_detection_test_loader(self.cfg, self.cfg.DATASETS.TEST[0], evaluators=[evaluator])\n",
    "        val_results = self._trainer.test(self.cfg, self.model, evaluators=[evaluator])[0]\n",
    "        val_loss = val_results[\"total_loss\"]\n",
    "        return val_loss\n",
    "\n",
    "    def save_checkpoint(self):\n",
    "        checkpointer = DetectionCheckpointer(self.model)\n",
    "        checkpointer.save(self.model_checkpoint_path)\n",
    "        \n",
    "\n",
    "def train_model(TRAIN,VALIDATION,folder):\n",
    "    cfg = get_cfg()\n",
    "    MODEL = 'faster_rcnn_X_101_32x8d_FPN_3x.yaml'\n",
    "    cfg.merge_from_file(model_zoo.get_config_file(\"COCO-Detection/\"+MODEL))\n",
    "    cfg.DATASETS.TRAIN = (TRAIN,)\n",
    "    cfg.DATASETS.TEST = (VALIDATION,)\n",
    "    cfg.DATALOADER.NUM_WORKERS = 2\n",
    "    #Uncomment if you want to use pre-trained weights for finetuning, not recommended for K fold training\n",
    "    # cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-Detection/\"+MODEL)  # Let training initialize from model zoo\n",
    "    \n",
    "    \n",
    "    cfg.SOLVER.IMS_PER_BATCH = 2  # This is the real \"batch size\" commonly known to deep learning people\n",
    "    cfg.SOLVER.BASE_LR = 0.004  # pick a good LR\n",
    "    cfg.SOLVER.STEPS = []        # milestones where LR is reduced, in this case there's no decay\n",
    "    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # The \"RoIHead batch size\". \n",
    "    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 80  \n",
    "    cfg.TEST.EVAL_PERIOD = 15000\n",
    "    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
    "    trainer = Early_stopping(cfg, early_stop_patience=5, model_checkpoint_path=\"model_checkpoint.pth\")\n",
    "    # Specify evaluators during testing\n",
    "    evaluator = COCOEvaluator(cfg.DATASETS.TEST[0], cfg, True, output_dir=\"./output/\")\n",
    "    trainer.resume_or_load(resume=False)\n",
    "    trainer.test(cfg, trainer.model, evaluators=[evaluator])\n",
    "    trainer.resume_or_load(resume=False)\n",
    "    trainer.train();\n",
    "    return cfg\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the results\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The given code block implements a function \"format_detectron2_predictions\" that converts the output of Detectron2 to a format that can be used by Cleanlab for identifying label errors. The function accepts the predicted instances and the number of classes as inputs. It processes the predicted bounding boxes and scores for each instance, and outputs a list of numpy arrays containing the bounding boxes and scores for each class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def format_detectron2_predictions(instances, num_classes):\n",
    "    \"\"\"\n",
    "    Format Detectron2 predictions into a list of NumPy arrays.\n",
    "\n",
    "    Args:\n",
    "        instances (detectron2.structures.Instances): Predicted instances from Detectron2.\n",
    "        num_classes (int): Number of classes.\n",
    "\n",
    "    Returns:\n",
    "        list of numpy.ndarray: Formatted predictions where each element represents a class and contains \n",
    "        an array of bounding box coordinates and scores.\n",
    "\n",
    "    \"\"\"\n",
    "    fields = instances.get_fields()\n",
    "    boxes = fields['pred_boxes'].tensor.numpy()\n",
    "    results = [[] for _ in range(num_classes)]\n",
    "\n",
    "    for i in range(len(fields['pred_classes'])):\n",
    "        pred_class = fields['pred_classes'][i].item()\n",
    "        scores = instances.get_fields()['scores'][i].item()\n",
    "        box_coordinates = list(boxes[i])\n",
    "        box_coordinates.append(scores)\n",
    "        results[pred_class].append(box_coordinates)\n",
    "\n",
    "    formatted_results = []\n",
    "    for i in results:\n",
    "        if len(i) == 0:\n",
    "            formatted_array = np.array(i, dtype=np.float32).reshape((0, 5))\n",
    "        else:\n",
    "            formatted_array = np.array(i, dtype=np.float32)\n",
    "        formatted_results.append(formatted_array)\n",
    "\n",
    "    return formatted_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for k in range(0,NUM_FOLDS):\n",
    "    result_dict = {}\n",
    "    train_data = pairs[k][0]\n",
    "    val_data = pairs[k][1]\n",
    "    cfg = train_model(train_data,val_data,\"COCO_TRAIN_\"+str(k)+\"_FOLD\")\n",
    "    evaluator = COCOEvaluator(val_data, output_dir=\"output\")\n",
    "    val_loader = build_detection_test_loader(cfg, val_data)\n",
    "    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")  # path to the model we just trained\n",
    "    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.1   # set a custom testing threshold\n",
    "    predictor = DefaultPredictor(cfg)\n",
    "    dataset = json.load(open(pairs[k][1],'rb'))\n",
    "    for image in dataset['images']:\n",
    "        im_name = os.path.join(TRAIN_PATH, image['file_name'])\n",
    "        im = cv2.imread(im_name)\n",
    "        outputs = predictor(im)\n",
    "        result_dict[im_name] = (format_detectron2_predictions(outputs[\"instances\"].to(\"cpu\"),cfg.MODEL.ROI_HEADS.NUM_CLASSES))\n",
    "    pickle.dump(result_dict,open(\"results_fold_\"+str(k)+\".pkl\",'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_dict = {}\n",
    "for k in range(0,NUM_FOLDS):\n",
    "    res_d = pickle.load(open(\"results_fold_\"+str(k)+'.pkl','rb'))\n",
    "    for r in res_d:\n",
    "        result_dict[r] = res_d[r]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = pickle.load(open(\"TRAIN_COCO_ALL_labels.pkl\",'rb'))\n",
    "results = []\n",
    "for i in dataset:\n",
    "    im_name = os.path.join(TRAIN_PATH, i['seg_map'].replace(\".png\",'.jpg'))\n",
    "    results.append(result_dict[im_name])\n",
    "pickle.dump(results,open(\"results_train_ALL.pkl\",'wb'))"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}