{"cells":[{"metadata":{},"cell_type":"markdown","source":"# [Global Wheat Detection](https://www.kaggle.com/c/global-wheat-detection)\n\nIn this notebook I will try to detect the wheat-heads by training the Faster R-CNN model implemented with PyTorch. I will try to keep this notebook clean, concise & well commented for everyone (including beginners) to understand. \n\n\n\nSo, let's begin the fun... \n\n![](https://i.imgflip.com/12aka8.jpg)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Libraries","execution_count":null},{"metadata":{"_uuid":"ddaffa4a-8183-4c8c-8313-7f6f0aeac8cb","_cell_guid":"28e5669f-7419-4aad-b560-97e77ea71763","trusted":true},"cell_type":"code","source":"#Generic Packages\nimport numpy as np\nimport os\nimport pandas as pd\nimport re\nfrom PIL import Image\nfrom IPython.display import display\nimport gc\nfrom PIL import Image\nimport warnings \nwarnings.filterwarnings(\"ignore\")\n\n#Data Augmentation Libraries\nimport albumentations as al\nfrom albumentations.pytorch.transforms import ToTensorV2\n\n#PyTorch Libraries\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\n#openCV\nimport cv2  \n\n#Plotting Libraries\nimport seaborn as sns; sns.set(font_scale=1.4)\nimport matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\nimport matplotlib.patches as patches\n\npd.plotting.register_matplotlib_converters()\n%matplotlib inline\nplt.style.use('seaborn-whitegrid')\npd.set_option('display.max_columns', 500)\n\nfrom tqdm import tqdm\nfrom pathlib import Path","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# Custom Class & Functions\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\nclass WheatDS(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        \n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # there is only one class\n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        # target['masks'] = None\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\nclass Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Load & Split","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"#Define Directories\ninput_dir = '../input/global-wheat-detection'\ntrain_dir = f'{input_dir}/train'\ntest_dir = f'{input_dir}/test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Train Image Number\nimg_number = -1 \n\n#Load Training Data\ntrain_data = pd.read_csv(f'{input_dir}/train.csv')\n\n#Add New Columns\ntrain_data['x'] = -1\ntrain_data['y'] = -1\ntrain_data['w'] = -1\ntrain_data['h'] = -1\n\n#Expand BBox\ntrain_data[['x', 'y', 'w', 'h']] = np.stack(train_data['bbox'].apply(lambda x: expand_bbox(x)))\n\n#Drop bbox column\ntrain_data.drop(columns=['bbox'], inplace=True)\n\n#Convert the data type for new columns\ntrain_data['x'] = train_data['x'].astype(np.float)\ntrain_data['y'] = train_data['y'].astype(np.float)\ntrain_data['w'] = train_data['w'].astype(np.float)\ntrain_data['h'] = train_data['h'].astype(np.float)\n\nimage_ids = train_data['image_id'].unique()\nvalid_ids = image_ids[img_number:]\ntrain_ids = image_ids[:img_number]\n\ntest_data = train_data[train_data['image_id'].isin(valid_ids)]\ntrain_data = train_data[train_data['image_id'].isin(train_ids)]\n\nprint(\"Validation Dataset Shape: \", test_data.shape, \"-- \", \"Training Dataset Shape: \",train_data.shape)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# EDA - Numerical","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"#Inspect\ntrain_data.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Data Types\ntrain_data.dtypes","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Images per Source\ntrain_data.groupby(\"source\").image_id.count()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Total Images in Training :', len(os.listdir(train_dir)))\nprint('Total Images in Test:', len(os.listdir(test_dir)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Unique Images\nprint(\"Unique Training Images: \", train_data['image_id'].nunique())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# EDA - Visual","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"#Visualize the Image Distribution across the different Sources\nplt.rcParams['figure.figsize'] = (10, 8)\nsns.countplot(train_data['source'], palette='pastel')\nplt.title('Image Distribution Across the Different Sources', fontsize= 15)\nplt.legend()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# Function to view images from dataset\ndef view_images(image_ids):\n    col = 5\n    row = min(len(image_ids) // col, 5)\n    fig, ax = plt.subplots(row, col, figsize=(16, 8))\n    fig.suptitle('Random Images of Wheat Heads from Training Dataset', fontsize=15)\n    ax = ax.flatten()\n    \n\n    for i, image_id in enumerate(image_ids):\n        image = cv2.imread('../input/global-wheat-detection/train' + '/{}.jpg'.format(image_id))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        ax[i].set_axis_off()\n        ax[i].imshow(image)\n\n#Running the function\nview_images(train_data.sample(n=10)['image_id'].values)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Simple Object Detection ","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# Function to read bbox and plot the image with the detected wheat heads\n\ndef get_bboxes(df, image_id):\n    image_bboxes = df[df.image_id == image_id]\n    \n    bboxes = []\n    for _,row in image_bboxes.iterrows():\n        bboxes.append((row.x, row.y, row.w, row.h))\n        \n    return bboxes\n\ndef plot_image(df, rows=3, cols=3, title='Wheat Head Detection in Training Images'):\n    fig, axs = plt.subplots(rows, cols, figsize=(10,10))\n    for row in range(rows):\n        for col in range(cols):\n            idx = np.random.randint(len(df), size=1)[0]\n            img_id = df.iloc[idx].image_id\n            \n            img = Image.open(train_dir + '/' + img_id + '.jpg')\n            axs[row, col].imshow(img)\n            \n            bboxes = get_bboxes(df, img_id)\n            \n            for bbox in bboxes:\n                rect = patches.Rectangle((bbox[0],bbox[1]),bbox[2],bbox[3],linewidth=1,edgecolor='r',facecolor='none')\n                axs[row, col].add_patch(rect)\n            \n            axs[row, col].axis('off')\n            \n    plt.suptitle(title)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Plotting the training image with detected wheat heads \nplot_image(train_data)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Build Model","execution_count":null},{"metadata":{"_uuid":"c3d93a46-fe4a-4bbb-b0a0-423e2d8da7a3","_cell_guid":"929d6c6c-9ba3-4a63-99df-a5f7de06b5f1","trusted":true},"cell_type":"code","source":"# load a model; pre-trained on COCO\n\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained_backbone=False)\nmodel.load_state_dict(torch.load(\"../input/torchvisionfasterrcnn/fasterrcnn_resnet50_fpn.pth\"))\n\nnum_classes = 2  # 1 class (wheat) + background\n\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data Augmentation","execution_count":null},{"metadata":{"_uuid":"1eb173f4-ea66-42c6-95f7-2d60341d3ebb","_cell_guid":"fa5cb5d7-7634-4676-82da-312c94db0937","trusted":true},"cell_type":"code","source":"# Albumentations\ndef get_train_transform():\n    return al.Compose([\n        al.RandomSizedBBoxSafeCrop(512, 512, erosion_rate=0.0, interpolation=1, p=1.0),\n        al.HorizontalFlip(p=0.5),\n        al.VerticalFlip(p=0.5),\n        al.RandomContrast(p=1.0),\n        al.RandomGamma(p=1.0),\n        al.RandomBrightness(p=1.0),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return al.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset Creation","execution_count":null},{"metadata":{"_uuid":"8733d23d-14ea-4eed-ac75-f19d7d6f4abf","_cell_guid":"0c78fe71-284d-4379-83e8-082e11ff235d","trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = WheatDS(train_data, train_dir, get_train_transform())\ntest_dataset = WheatDS(test_data, train_dir, get_valid_transform())\n\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=16,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\ntest_data_loader = DataLoader(\n    test_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nmodel.to(device)\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)\n#lr_scheduler = None\n\nnum_epochs = 9\n\nloss_hist = Averager()\nitr = 1000","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train the Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"for epoch in range(num_epochs):\n    loss_hist.reset()\n    \n    for images, targets, image_ids in train_data_loader:\n        \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\n        loss_dict = model(images, targets)\n\n        losses = sum(loss for loss in loss_dict.values())\n        loss_value = losses.item()\n\n        loss_hist.send(loss_value)\n\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n\n        if itr % 25 == 0:\n            print(f\"Iteration #{itr} loss: {loss_value}\")\n\n        itr += 1\n    \n    # update the learning rate\n    if lr_scheduler is not None:\n        lr_scheduler.step()\n\n    print(f\"Epoch #{epoch} loss: {loss_hist.value}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Submission","execution_count":null},{"metadata":{"_uuid":"433cc43a-6fa8-47c6-95ee-59f65690fb10","_cell_guid":"203e871c-1c2f-4f6c-8068-94301f36b0df","trusted":true},"cell_type":"code","source":"#Function for making predictions\ndef prediction(image_path,model,device):\n    model.eval()\n    image = cv2.imread(image_path, cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n    image /= 255.0\n    images = torch.from_numpy(image).float().permute(2,0,1).unsqueeze(0).to(device)\n    outputs = model(images)\n\n\n    outputs = [{k: v.detach().cpu().numpy() for k, v in t.items()} for t in outputs]\n    boxes = outputs[0][\"boxes\"]\n    scores = outputs[0][\"scores\"]\n    valid_boxes = boxes[scores > 0.5]\n    valid_scores = scores[scores > 0.5]\n    return valid_boxes, valid_scores","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"0b9cf11d-a6b9-47ca-816c-4952c8b801e0","_cell_guid":"ad8472f3-83d3-4a39-93ba-fb1b005363d6","trusted":true},"cell_type":"code","source":"#Inspecting the sample submission\nsubmission = pd.read_csv(f'{input_dir}/sample_submission.csv')\nsubmission.head()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"9fb694ca-b0f5-4627-968f-572d12c217d8","_cell_guid":"a07c23cf-affe-4dfe-ac9f-c51425ecd968","trusted":true},"cell_type":"code","source":"sub_dir = '../input/global-wheat-detection/test'\n\nsubmission = pd.read_csv(f'{input_dir}/sample_submission.csv')\n\n\n\nroot_image = Path(test_dir)\ntest_images = [root_image / f\"{img}.jpg\" for img in submission.image_id]\n\n\nsubmission = []\nmodel.eval()\n\nfor image in tqdm(test_images):\n    boxes, scores = prediction(str(image),model,device)\n    prediction_string = []\n    for (x_min,y_min,x_max,y_max),s in zip(boxes,scores):\n        x = round(x_min)\n        y = round(y_min)\n        h = round(x_max-x_min)\n        w = round(y_max-y_min)\n        prediction_string.append(f\"{s} {x} {y} {h} {w}\")\n    prediction_string = \" \".join(prediction_string)\n    \n    submission.append([image.name[:-4],prediction_string])\n\nsample_submission = pd.DataFrame(submission, columns=[\"image_id\",\"PredictionString\"])\nsample_submission.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Please upvote if you find it helpful :-)","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}