{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Download TorchVision repo to use some files from\n# references/detection\n!git clone https://github.com/pytorch/vision.git\n\n!cp /kaggle/working/vision/references/detection/utils.py /kaggle/working/\n!cp /kaggle/working/vision/references/detection/transforms.py /kaggle/working/\n!cp /kaggle/working/vision/references/detection/coco_eval.py /kaggle/working/\n!cp /kaggle/working/vision/references/detection/engine.py /kaggle/working/\n!cp /kaggle/working/vision/references/detection/coco_utils.py /kaggle/working/\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport torch.nn\nimport pandas\nimport re\nimport os\nimport zipfile\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom torch.utils.data import TensorDataset, DataLoader\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\nfrom PIL import Image\nfrom engine import train_one_epoch, evaluate\nimport transforms as T\nimport utils\nimport matplotlib.patches as patches","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR = \"/kaggle/input/global-wheat-detection/\"\npd_train = pandas.read_csv(DIR+\"train.csv\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **COLLECTING DATA TO ARRAY**","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"new_bbox = []\n\nimage_id = pd_train[\"image_id\"].to_numpy()\nsource = pd_train['source'].to_numpy()\nbbox = pd_train[\"bbox\"].to_numpy()\n\nfor b in bbox:\n    #save new bbox\n    box = b.strip(\"[]\").split(\",\")\n    x,y,w,h = np.array(box).astype(float)\n    new_bbox.append([float(box[0]),float(box[1]),float(box[2]),float(box[3])])\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class dictImage():\n    def __init__(self, DIR, image_id, bbox):\n        self.DIR = DIR\n        self.image_id = image_id\n        self.bbox = bbox #like a coordinate from image\n        self.dict_image = {}\n        self.list_data = []\n    \n    def buildDict(self):\n        self.dict_image = {}\n        idx = 0\n        for image_id, bbox in zip(self.image_id, self.bbox):\n            path = self.DIR+\"/train/\"+ image_id+\".jpg\"\n            \n            if image_id not in self.dict_image:\n                self.dict_image[image_id] = {}\n                self.dict_image[image_id][\"path\"] =  path\n                self.dict_image[image_id][\"bbox\"] = [bbox]\n            else:\n                self.dict_image[image_id][\"bbox\"].append(bbox)\n\n    def convertToListData(self):\n        self.list_data = []\n        for key, val in self.dict_image.items():\n            self.list_data.append(val)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_train = dictImage(DIR, image_id, new_bbox)\ndata_train.buildDict()\ndata_train.convertToListData()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(data_train.list_data)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# BUILD DATASET","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transform(train):\n    transforms = []\n    # converts the image, a PIL image, into a PyTorch Tensor\n    transforms.append(T.ToTensor())\n    if train:\n        # during training, randomly flip the training images\n        # and ground-truth for data augmentation\n        transforms.append(T.RandomHorizontalFlip(0.5))\n    return T.Compose(transforms)\ntransf_train = get_transform(train=True)\ntransf_test = get_transform(train=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BuildDataset(torch.utils.data.Dataset):\n    def __init__(self, list_data, transforms=None):\n        self.transforms = transforms\n        self.list_data = list_data\n        self.reduce_dim = 0.25\n    \n    def getMasks(self, path, bbox):\n        sv = True\n        img = cv2.imread(path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n        nw = int(self.reduce_dim * img.shape[0])\n        nh = int(self.reduce_dim * img.shape[1])\n        img = cv2.resize(img,(nw,nh)) #reduce dimention of image for reduce time consuming\n        \n        masks = np.zeros_like(img)\n        x,y,w,h =  bbox\n        x1 = int(x * self.reduce_dim)\n        x2 = int((x+w) * self.reduce_dim)\n        y1 = int(y * self.reduce_dim)\n        y2 = int((y+h) * self.reduce_dim)\n        \n        masks[y1:y2, x1:x2] = 1 #remember that arrays always start from row to column or y to x\n        \n        if(x1 == x2):\n            sv = False\n        if(y1 == y2):\n            sv = False\n        \n        xy = [x1,y1,x2,y2]\n        \n        return masks, xy, sv\n\n    def __getitem__(self, idx):\n        data = self.list_data[idx]\n        path = data['path']\n        \n        #reduce image dimention to reduce time processing\n        img = Image.open(path).convert(\"RGB\")\n        nw = int(self.reduce_dim * img.size[0])#new width\n        nh = int(self.reduce_dim * img.size[1])#new height\n        img = img.resize((nw,nh))#new dimention image\n        \n        masks = []\n        n_bbox = []\n        lbl = []\n        \n        for bbox in data['bbox']:\n            \n            bbox = np.array(bbox).astype(int)\n            m, box, sv = self.getMasks(path, bbox) #sv variable to validate data whether it is used or not\n            if sv == True:\n                lbl.append(1)\n                masks.append(m)\n                n_bbox.append(box)\n                \n        target = {}\n        target[\"boxes\"] = torch.as_tensor(n_bbox, dtype=torch.float32)\n        target[\"labels\"] = torch.as_tensor(lbl, dtype=torch.int64)\n        target[\"masks\"] = torch.as_tensor(masks, dtype=torch.uint8)\n        \n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n\n        return img, target\n    \n    def __len__(self):\n        return len(self.list_data)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.manual_seed(1)\ndataset = BuildDataset(data_train.list_data, transf_train)\n\n# define training data loaders\ndata_loader = torch.utils.data.DataLoader(\n    dataset, batch_size=5, shuffle=True, num_workers=6,\n    collate_fn=utils.collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **DEFINE MODEL**","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_instance_segmentation_model(num_classes):\n    # load an instance segmentation model pre-trained on COCO\n    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)\n\n    # get the number of input features for the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    # replace the pre-trained head with a new one\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    # now get the number of input features for the mask classifier\n    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n    hidden_layer = 256\n    # and replace the mask predictor with a new one\n    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,\n                                                       hidden_layer,\n                                                       num_classes)\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n# our dataset has two classes only - background and person\nnum_classes = 2#background and wheat\n\n# get the model using our helper function\nmodel = get_instance_segmentation_model(num_classes)\n# move model to the right device\nmodel.to(device)\n\n# construct an optimizer\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005,\n                            momentum=0.9, weight_decay=0.0005)\n\n# and a learning rate scheduler which decreases the learning rate by\n# 10x every 3 epochs\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,\n                                               step_size=3,\n                                               gamma=0.1)\nprint(\"device\",device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**LET'S TRAIN DATA**","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"I only did training for 2 epochs. I have previously tried training 5 epochs but I did not have enough time because I was limited to 9 hours for this assignment","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# let's train it for 10 epochs\nnum_epochs = 2\n\nfor epoch in range(num_epochs):\n    # train for one epoch, printing every 100 iterations\n    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=20)\n    # update the learning rate\n    lr_scheduler.step()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def showImage(img, boxes, scores):\n    img = img.mul(255) \n    img = img.permute(1,2,0).byte().cpu().numpy()\n    fig,ax = plt.subplots(1)\n    fig.set_figheight(8)\n    fig.set_figwidth(8)\n    ax.imshow(img)\n    \n    for box, score in zip(boxes, scores):\n        if score >= 0.5:\n            x1,y1,x2,y2= box\n            x = x1 \n            y = y1\n            w = x2 - x1\n            h = y2-y1\n            rect = patches.Rectangle((x,y),w,h,linewidth=1,edgecolor='r',facecolor='none')\n            ax.add_patch(rect)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\ndir_test = DIR+\"test/\"\n\nlist_test = os.listdir(dir_test)\nfor f in list_test:\n    pth = dir_test + f\n    # pth = data_train.list_data[100]['path']\n    img, _ = transf_test(Image.open(pth),None)\n    img = img.unsqueeze(0).cuda()\n    predict = model(img)\n    boxes = predict[0]['boxes']\n    scores = predict[0]['scores']\n    img = img.squeeze(0)\n    showImage(img, boxes, scores)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}