{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom matplotlib import pyplot as plt\n\n#from tqdm import tqdm_notebook as tqdm\n#from tqdm import tqdm \nfrom tqdm.notebook import tqdm as tqdm\n\nimport cv2\nimport os\nimport re\n\nimport random\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\n\nimport ast\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","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"DATA_PATH = \"../input/global-wheat-detection/\"\nTRAIN_DIR = os.path.join(DATA_PATH, \"train\")\nTEST_DIR = os.path.join(DATA_PATH, \"test\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv(os.path.join(DATA_PATH, \"train.csv\"))\ndf.head(5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Shape of Dataframe\nprint(f\"Shape of train DataFrame:{df.shape}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Extract x,y,w,h from bbox\ndef extract_bbox(DataFrame):\n    DataFrame[\"x\"] = [np.float(ast.literal_eval(i)[0]) for i in DataFrame[\"bbox\"]]\n    DataFrame[\"y\"] = [np.float(ast.literal_eval(i)[1]) for i in DataFrame[\"bbox\"]]\n    DataFrame[\"w\"] = [np.float(ast.literal_eval(i)[2]) for i in DataFrame[\"bbox\"]]\n    DataFrame[\"h\"] = [np.float(ast.literal_eval(i)[3]) for i in DataFrame[\"bbox\"]]\n    \nextract_bbox(df)\ndf.head()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training and validation split\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_split = 0.8\nimages_id   = df[\"image_id\"].unique() \ntrain_ids   = images_id[:int(len(images_id)*train_split)]\nvalid_ids   = images_id[int(len(images_id)*train_split):]\n\nprint(f'Total Images Number: {len(images_id)}')\nprint(f'Number of training images: {len(train_ids)}')\nprint(f'Number of Valid images: {len(valid_ids)}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = df[df[\"image_id\"].isin(train_ids)]\nvalid_df = df[df[\"image_id\"].isin(valid_ids)]\n\nprint(f'Shape of train_df: {train_df.shape}')\nprint(f'Shape of valid_df: {valid_df.shape}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Data Transform - Albumentation\ndef get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n    def __init__(self, dataframe, image_dir, transform=None):\n        super().__init__()\n        self.dataframe = dataframe\n        self.image_dir = image_dir\n        self.transform = transform\n        self.image_ids = dataframe[\"image_id\"].unique()\n        \n    def __getitem__(self, idx):\n        image_id = self.image_ids[idx]\n        details = self.dataframe[self.dataframe[\"image_id\"]==image_id]\n        img_path = os.path.join(TRAIN_DIR, image_id)+\".jpg\"\n        image = cv2.imread(img_path, cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        \n        #Row of Dataframe of a particular index.\n        boxes = details[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        \n        #To find area\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        \n        #COnvert it into tensor dataType\n        area = torch.as_tensor(area, dtype=torch.float32)\n        \n        # there is only one class\n        labels = torch.ones((details.shape[0],), dtype=torch.int64)\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((details.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor(idx) ### <------------ New change list has been removed\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n        \n        if self.transform:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            \n            sample = self.transform(**sample)\n            image = sample['image']\n            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n            target[\"boxes\"] = torch.as_tensor(target[\"boxes\"], dtype=torch.long)\n        \n        return image, target     #, image_id\n    \n    def __len__(self) -> int:\n        return len(self.image_ids)\n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = WheatDataset(train_df, TRAIN_DIR, get_train_transform())\nvalid_dataset = WheatDataset(valid_df, TRAIN_DIR, get_valid_transform())\n\nprint(f\"Length of train_dataset: {len(train_dataset)}\")\nprint(f\"Length of test_dataset: {len(valid_dataset)}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"indices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_images(n_num, random_selection=True):\n\n    if random_selection:\n        index = random.sample(range(0, len(train_df[\"image_id\"].unique())-1), n_num)\n    else:\n        index = range(0, n_num)\n    plt.figure(figsize=(15,15))\n    fig_no = 1\n    \n    for i in index:\n        images, targets = train_dataset.__getitem__(i)\n        sample = np.array(np.transpose(images, (1,2,0)))\n        boxes = targets[\"boxes\"].numpy().astype(np.int32)\n    \n        #Plot figure/image\n\n        for box in boxes:\n            cv2.rectangle(sample,(box[0], box[1]),(box[2], box[3]),(255,223,0), 2)\n        plt.subplot(n_num/2, n_num/2, fig_no)\n        plt.imshow(sample)\n        fig_no+=1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_images(4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained = True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_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":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\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)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = None\n\nnum_epochs = 8","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import time\n\nitr=1\n\ntotal_train_loss = []\ntotal_valid_loss = []\n\nlosses_value = 0\nfor epoch in range(num_epochs):\n  \n    start_time = time.time()\n    train_loss = []\n    model.train()\n    \n    #<-----------Training Loop---------------------------->\n    pbar = tqdm(train_data_loader, desc = 'description')\n    for images, targets in pbar:\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        losses_value = losses.item()\n        train_loss.append(losses_value)        \n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n        \n        pbar.set_description(f\"Epoch: {epoch+1}, Batch: {itr}, loss: {losses_value}\")\n        itr+=1\n\n    epoch_train_loss = np.mean(train_loss)\n    total_train_loss.append(epoch_train_loss)\n    \n    \n    #<---------------Validation Loop---------------------->\n    with torch.no_grad():\n        valid_loss = []\n\n        for images, targets in valid_data_loader:\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            # If you need validation losses\n            model.train()\n            # Calculate validation losses\n            loss_dict = model(images, targets)\n            losses = sum(loss for loss in loss_dict.values())\n            loss_value = losses.item()\n            valid_loss.append(loss_value)\n            \n    epoch_valid_loss = np.mean(valid_loss)\n    total_valid_loss.append(epoch_valid_loss)\n    \n    print(f\"Epoch Completed: {epoch+1}/{num_epochs}, Time: {time.time()-start_time},\\\n    Train Loss: {epoch_train_loss}, Valid Loss: {epoch_valid_loss}\")    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import seaborn as sns\nplt.figure(figsize=(8,5))\nsns.set_style(style=\"whitegrid\")\nsns.lineplot(range(1, len(total_train_loss)+1), total_train_loss, label=\"Training Loss\")\nsns.lineplot(range(1, len(total_train_loss)+1), total_valid_loss, label=\"Valid Loss\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Loss\")\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), 'fasterrcnn_best_resnet50.pth')","execution_count":null,"outputs":[]}],"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}