{"cells":[{"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 torch\nfrom torchvision import models,transforms\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\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\nfrom matplotlib import pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!git clone https://github.com/Paperspace/DataAugmentationForObjectDetection.git\n%cd /kaggle/working/DataAugmentationForObjectDetection\nfrom data_aug.data_aug import *\nfrom data_aug.bbox_util import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv('/kaggle/input/global-wheat-detection/train.csv')\nprint(train_df.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img_ids = train_df['image_id'].unique()\nvalid_ids = img_ids[-100:]\ntrain_ids = img_ids[:-100]\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\nvalid_df.shape,train_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## New transform function for image augmentation","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Sequence(object):\n    def __init__(self, augmentations, probs = 0.5):\n        self.augmentations = augmentations\n        self.probs = probs\n    def __call__(self, image, bboxes):\n#         bboxes = target[\"boxes\"]\n        for i, augmentation in enumerate(self.augmentations):\n            if type(self.probs) == list:\n                prob = self.probs[i]\n            else:\n                prob = self.probs\n            img = image.copy()\n            boxes = bboxes.copy()\n            if random.random() < prob:\n                img, boxes = augmentation(img, boxes)\n        return img, boxes\n# t1 = Sequence([RandomHorizontalFlip(1), RandomScale(0.2, diff = True), RandomRotate(10)])\nt2 = Sequence([RandomHSV(40, 40, 30),RandomHorizontalFlip(), RandomScale(), RandomTranslate(), RandomRotate(10), RandomShear()])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GWDataset(torch.utils.data.Dataset):\n    def __init__(self, dataframe, image_dir, train=False):\n        super().__init__()\n\n        self.img_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.train = train\n\n    def __getitem__(self, index):\n        img_idx = self.img_ids[index]\n        img_name = str(img_idx+'.jpg')\n        # load images ad masks\n        img_path = os.path.join(self.image_dir, img_name)\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).astype(np.float32)\n       # get bounding box coordinates for each mask\n        num_bbxs = len(self.df[self.df['image_id']==img_idx])\n        bbxs = self.df[self.df['image_id']==img_idx]\n        boxes = []\n        area = []\n#         print(bbxs)\n        for t in range(num_bbxs):\n            l = bbxs.iloc[t]['bbox'].split(',')\n#             print(l)\n            xmin,ymin,w,h = float(l[0][1:]),float(l[1][1:]),float(l[2][1:]),float(l[3][1:-1])\n            xmax = xmin+w\n            ymax = ymin+h\n            area.append(w*h)\n            boxes.append([xmin, ymin, xmax, ymax])\n\n        # there is only one class\n        labels = torch.ones((num_bbxs,), dtype=torch.int64)\n\n        imag_id = torch.tensor([index])\n        # suppose all instances are not crowd\n        area = torch.as_tensor(area, dtype=torch.float32)\n        iscrowd = torch.zeros((num_bbxs,), dtype=torch.int64)\n#         print(boxes)\n        target = {}\n#         target[\"boxes\"] = boxes\n        target[\"labels\"] = labels\n        target[\"image_id\"] = imag_id\n        target[\"area\"] = area\n        target[\"iscrowd\"] = iscrowd\n        imge = img.copy()\n        if self.train is True:\n            imge,boxes = t2(imge, np.array(boxes))\n        imge = transforms.Compose([transforms.ToPILImage() ,transforms.ColorJitter(hue=0.3),transforms.ToTensor()])(np.uint8(imge))\n#         imge = imge/255\n#         imge = transforms.ToTensor()(imge)\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        target[\"boxes\"] = boxes\n        return imge, target\n\n    def __len__(self):\n        return (self.img_ids.shape[0])\ndataset = GWDataset(valid_df,'/kaggle/input/global-wheat-detection/train/',train = True)\ndataset[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(4):\n    a = dataset[i][0].permute(1,2,0)\n    b = dataset[i+1][0].permute(1,2,0)\n    fig, (ax1, ax2) = plt.subplots(1, 2)\n    ax1.imshow(a)\n    ax2.imshow(b)\n    i=i+1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"##  CHOICE BETWEEN MODEL PRETRAINED ON image_net vs Global Wheat Detection Challenge\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n# model = torch.load(\"/kaggle/input/gwd-model/fasterrcnn_resnet50_fpn.pth\",map_location='cpu')","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":"class 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\n        \ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\n# split the dataset in train and test set\n\nDIR_TRAIN = '/kaggle/input/global-wheat-detection/train/'\ntrain_dataset = GWDataset(train_df, DIR_TRAIN,True)\nvalid_dataset = GWDataset(valid_df, DIR_TRAIN, False)\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=8,\n    shuffle=True,\n    num_workers=0,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=4,\n    shuffle=False,\n    num_workers=0,\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]\n\noptimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.train()\nmodel.to(device)#,dtype=torch.float)\nnum_epochs = 15\n\nloss_hist = Averager()\nitr = 1\n\nfor epoch in range(num_epochs):\n    loss_hist.reset()\n    \n    for images, targets in train_data_loader:\n#         try:\n#             print(\"!\")\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 % 50 == 0:\n                print(f\"Iteration #{itr} loss: {loss_value}\")\n\n            itr += 1\n#         except:\n#             print(\"F\",itr)\n#             pass\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":{"trusted":true},"cell_type":"code","source":"# pick one image from the test set\nfor i in range(1):\n    img, _ = valid_dataset[i]\n    # put the model in evaluation mode\n    model.eval()\n    with torch.no_grad():\n        prediction = model([img.to(device,dtype = torch.float32)])\n    sample = valid_dataset[i][0].permute(1,2,0).numpy()\n    boxes = prediction[0]['boxes'].cpu().numpy().astype(np.int32)\n    # boxe = boxes.reshape((4,-1))\n    scores = prediction[0]['scores'].cpu().numpy()\n    im = np.array(img.permute(1,2,0))\n    # plt.imshow(sample)\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\n    color = (220,0,0)\n    for i in range(len(boxes)):\n    #     print(boxes[i])\n        if scores[i]>0.90:\n            cv2.rectangle(im,(int(boxes[i][0]), int(boxes[i][1])),(int(boxes[i][2]), int(boxes[i][3])),color, 5)\n    ax.set_axis_off()\n    ax.imshow(im)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dir_test = \"/kaggle/input/global-wheat-detection/test\"\npreprocess = transforms.Compose([transforms.ToTensor()])\ncolor = (220,0,0)\nresults = []\nfor img_file in os.listdir(dir_test):\n    result = []\n    img = cv2.imread(os.path.join(dir_test,img_file))\n    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).astype(np.float32)\n    img = img/255.0\n    img_t = preprocess(img)\n    pred = model([img_t.to(device,dtype = torch.float32)])\n    bboxes = pred[0]['boxes'].cpu().detach().numpy()\n    bscores = pred[0]['scores'].cpu().detach().numpy()\n    img_name = img_file.split('.')[:-1]\n    for i in range(len(bboxes)):\n        if bscores[i]>0.5:\n            result.append((bscores[i],bboxes[i]))\n    results.append((str(img_name[0]),result))\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n#     for box in bboxes:\n#         cv2.rectangle(img,(int(box[0]), int(box[1])),(int(box[2]), int(box[3])),color, 5)\n    for i in result:\n        if i[0]>0.5:\n            box = i[1]\n            cv2.rectangle(img,(int(box[0]), int(box[1])),(int(box[2]), int(box[3])),color, 5)\n    plt.figure()\n    ax.set_axis_off()\n    ax.imshow(img)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model, '/kaggle/working/fasterrcnn_resnet50_fpn_new.pth')\ntorch.save(model.state_dict(), '/kaggle/working/fasterrcnn_resnet50_fpn_statedict.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# res = []\n# for result in results:\n# #     print(result[0],end='')\n#     pred_str = []\n#     for box in result[1]:\n#         pred_str.append(box[0])\n#         pred_str.append(box[1][0])\n#         pred_str.append(box[1][1])\n#         pred_str.append(box[1][2]-box[1][0])\n#         pred_str.append(box[1][3]-box[1][1])\n#     pred = {}\n#     pred['image_id'] = str(result[0])\n#     pred['PredictionString'] = ' '.join(str(i) for i in pred_str)\n#     res.append(pred)\n# test_df = pd.DataFrame(res, columns=['image_id', 'PredictionString'])\n# print(test_df)\n# test_df.to_csv(\"/kaggle/working/submission.csv\",index=False)\n# %cd /kaggle/working/\n# !cat submission.csv","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}