{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"from IPython.display import Image\nImage(\"/kaggle/input/global-wheat-detection/train/00ea5e5ee.jpg\", width=500)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\n### **Click [here](https://www.youtube.com/watch?v=RzHPiHr4fAM) for a video walkthrough**   \n\n* Finetunes a pre-trained FasterRCNN model.\n* This notebook requires **internet access**, to download the pre-trained network.\n* It has been tested on Kaggle as a kernel, with **GPU enabled**","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Notebook Overview\n- Data Pre-processing\n- Preparing a Pre-trained Bounding Box Detector\n- Preparing a Training Data Loader\n- Training\n- Validating the Training","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\n\nimport torch\n\nfrom torchvision.models.detection import fasterrcnn_resnet50_fpn\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n\nfrom torch.utils.data import DataLoader, Dataset\n\nfrom matplotlib import pyplot as plt\n# Set default figure size\nplt.rcParams['figure.figsize'] = (10.0, 10.0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"_","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Data Pre-processing","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define File Path Constants\nINPUT_DIR = os.path.abspath('/kaggle/input/global-wheat-detection')\nTRAIN_DIR = os.path.join(INPUT_DIR, \"train\")\n\n# Load and Show Training Labels\npd.read_csv(os.path.join(INPUT_DIR, \"train.csv\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Convenience functions for loading images\ndef read_image_from_path(image_path):\n    image = cv2.imread(image_path)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    return image\n\ndef read_image_from_train_folder(image_id):\n    path = os.path.join(TRAIN_DIR, image_id + \".jpg\")\n    return read_image_from_path(path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test loader functions\nsample_image_id = \"b6ab77fd7\"\nplt.imshow(read_image_from_train_folder(sample_image_id))\n_ = plt.title(sample_image_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Functions for parsing bounding box string into x1, y1, x2, y2\ndef parse_bbox_text(string_input):\n    input_without_brackets = re.sub(\"\\[|\\]\", \"\", string_input)\n    input_as_list = np.array(input_without_brackets.split(\",\"))\n    return input_as_list.astype(np.float) \n\ndef xywh_to_x1y1x2y2(x,y,w,h):\n    return np.array([x,y,x+w,y+h])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Parse training bounding box labels\ntrain_df = pd.read_csv(os.path.join(INPUT_DIR, \"train.csv\"))\nbbox_series = train_df.bbox.apply(parse_bbox_text)\n\nxywh_df = pd.DataFrame(bbox_series.to_list(), columns=[\"x\", \"y\", \"w\", \"h\"])\n\nx2_df = pd.DataFrame(xywh_df.x + xywh_df.w, columns=[\"x2\"])\ny2_df = pd.DataFrame(xywh_df.y + xywh_df.h, columns=[\"y2\"])\n\n# Update training dataframe with parsed labels\ntrain_df = train_df.join([xywh_df, x2_df, y2_df])\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Convenience function for drawing a list of bounding box coordinates on and image\ndef draw_boxes_on_image(boxes, image, color=(255,0,0)):    \n    for box in boxes:\n        cv2.rectangle(image,\n                      (int(box[0]), int(box[1]) ),\n                      (int(box[2]), int(box[3]) ),\n                      color, 3)\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Sample a random training instance and draw the labelled bounding boxes\nsample_image_id =  train_df.image_id.sample().item()\n\nsample_image = read_image_from_train_folder(sample_image_id)\nsample_bounding_boxes = train_df[train_df.image_id == sample_image_id][[\"x\", \"y\",\"x2\",\"y2\"]]\n\nplt.imshow(draw_boxes_on_image(sample_bounding_boxes.to_numpy(), sample_image, color=(0,200,200)))\n_ = plt.title(sample_image_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":" ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Preparing a Pre-trained Bounding Box Detector\n","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Download a pre-trained bounding box detector\nmodel = fasterrcnn_resnet50_fpn(pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Replace the pre-trained bounding box detector head with\n# a new one that predicts our desired 2 classes {BACKGROUND, WHEAT}\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_channels=in_features, num_classes=2)\n\n# Verify the model architecture\nmodel.roi_heads","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":" ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Preparing a Training Data Loader","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Determine device to run on. GPU is highly recommended\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\ndef move_batch_to_device(images, targets):\n    images = list(image.to(device) for image in images)\n    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n    return images, targets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## Split data into training and validation subsets\nunique_image_ids = train_df['image_id'].unique()\n\nn_validation = int(0.2 * len(unique_image_ids))\nvalid_ids = unique_image_ids[-n_validation:]\ntrain_ids = unique_image_ids[:-n_validation]\n\nvalidation_df = train_df[train_df['image_id'].isin(valid_ids)]\ntraining_df = train_df[train_df['image_id'].isin(train_ids)]\n\nprint(\"%i training samples\\n%i validation samples\" % (len(training_df.image_id.unique()),len(validation_df.image_id.unique())) )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Inherit from pytorch Dataset for convenience\n\nclass WheatDataset(Dataset):\n\n    def __init__(self, dataframe):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n\n    def __len__(self) -> int:\n        return len(self.image_ids)\n    \n    def __getitem__(self, index: int):\n\n        image_id = self.image_ids[index]\n        image = read_image_from_train_folder(image_id).astype(np.float32)\n        # Scale to [0,1] range expected by the pre-trained model\n        image /= 255.0\n        # Convert the shape from [h,w,c] to [c,h,w] as expected by pytorch\n        image = torch.from_numpy(image).permute(2,0,1)\n        \n        records = self.df[self.df['image_id'] == image_id]\n        \n        boxes = records[['x', 'y', 'x2', 'y2']].values\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        n_boxes = boxes.shape[0]\n        \n        # there is only one foreground class, WHEAT\n        labels = torch.ones((n_boxes,), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        \n        return image, target","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Create pytorch data loaders for training and validation\n\ntrain_dataset = WheatDataset(training_df)\nvalid_dataset = WheatDataset(validation_df)\n\n# A function to bring images with different\n# number of bounding boxes into the same batch\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\nis_training_on_cpu = device == torch.device('cpu')\nbatch_size = 4 if is_training_on_cpu else 16\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    shuffle=True,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test the data loader\nbatch_of_images, batch_of_targets = next(iter(train_data_loader))\n\nsample_boxes = batch_of_targets[0]['boxes'].cpu().numpy().astype(np.int32)\nsample_image = batch_of_images[0].permute(1,2,0).cpu().numpy() # convert back from pytorch format\n\nplt.imshow(draw_boxes_on_image(sample_boxes, sample_image, color=(0,200,200)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"_","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Set up the optimiser\noptimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_epochs = 1 if is_training_on_cpu else 3\n\n# Prepare the model for training\nmodel = model.to(device)\nmodel.train()\n    \nfor epoch in range(num_epochs):\n    print(\"Epoch %i/%i \" % (epoch + 1, num_epochs) )\n    average_loss = 0\n    for batch_id, (images, targets) in enumerate(train_data_loader):\n        # Prepare the batch data\n        images, targets = move_batch_to_device(images, targets)\n\n        # Calculate losses\n        loss_dict = model(images, targets)\n        batch_loss = sum(loss for loss in loss_dict.values()) / len(loss_dict)\n        \n        # Refresh accumulated optimiser state and minimise losses\n        optimizer.zero_grad()\n        batch_loss.backward()\n        optimizer.step()\n        \n        # Record stats\n        loss_value = batch_loss.item()\n        average_loss = average_loss + (loss_value - average_loss) / (batch_id + 1)\n        print(\"Mini-batch: %i/%i Loss: %.4f\" % ( batch_id + 1, len(train_data_loader), average_loss), end='\\r')\n        if batch_id % 100 == 0:\n            print(\"Mini-batch: %i/%i Loss: %.4f\" % ( batch_id + 1, len(train_data_loader), average_loss))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Validating the Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Prepare the model for inference\nmodel.eval()\n\n# Prepare a validation results generator\ndef make_validation_iter():\n    valid_data_iter = iter(valid_data_loader)\n    for images, targets in valid_data_iter:\n        images, targets = move_batch_to_device(images, targets)\n\n        cpu_device = torch.device(\"cpu\")\n        outputs = model(images)\n        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]\n        for image, output, target in zip(images, outputs, targets): \n            predicted_boxes = output['boxes'].cpu().detach().numpy().astype(np.int32)\n            ground_truth_boxes = target['boxes'].cpu().numpy().astype(np.int32)\n            image = image.permute(1,2,0).cpu().numpy()\n            yield image, ground_truth_boxes, predicted_boxes\n\nvalidation_iter = make_validation_iter()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image, ground_truth_boxes, predicted_boxes = next(validation_iter)\nimage = draw_boxes_on_image(predicted_boxes, image, (255,0,0))\nimage = draw_boxes_on_image(ground_truth_boxes, image , (0,255,0))\nplt.imshow(image)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":" ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Wrapping Up","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), 'fasterrcnn_gwd_finetuned.pth')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## [Link to inference kernel](https://www.kaggle.com/fulkast/v1-minimal-global-wheat-detection-inference/)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Other Useful Resources","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"* Fine-tuning pre-trained models  https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html   \n* Inheriting from Pytorch Dataset https://pytorch.org/docs/stable/data.html\n* FasterRCNN.                     https://arxiv.org/abs/1506.01497 \n* Why Momentum Really Works              https://distill.pub/2017/momentum/","execution_count":null}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}