{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport torch\nimport os,pdb\nimport time\nimport shutil\nimport torch.nn as nn\nfrom skimage import io\nfrom PIL import Image\nimport torchvision\nimport cv2\nfrom tqdm import tqdm\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\nfrom albumentations.pytorch import ToTensor\nfrom torchvision import utils\nfrom albumentations import (HorizontalFlip, ShiftScaleRotate, VerticalFlip, Normalize,Flip,\n                            Compose, GaussNoise)\n\n\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:20:09.103648Z","iopub.execute_input":"2021-06-03T12:20:09.104022Z","iopub.status.idle":"2021-06-03T12:20:12.143022Z","shell.execute_reply.started":"2021-06-03T12:20:09.103941Z","shell.execute_reply":"2021-06-03T12:20:12.142134Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''\nhelp functions\n'''\ndef get_transforms(phase):\n            list_transforms = []\n            if phase == 'train':\n                list_transforms.extend([\n                       ToTensor()\n                         ])\n            list_trfms = Compose(list_transforms,\n                                 bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n            return list_trfms\n        \n# batching\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndef save_ckp(state, is_best, checkpoint_path, best_model_path):\n    \"\"\"\n    state: checkpoint we want to save\n    is_best: is this the best checkpoint; min validation loss\n    checkpoint_path: path to save checkpoint\n    best_model_path: path to save best model\n    \"\"\"\n    # save checkpoint data to the path given, checkpoint_path\n    torch.save(state, checkpoint_path)\n    # if it is a best model, min validation loss\n    if is_best:\n        # copy that checkpoint file to best path given, best_model_path\n        shutil.copyfile(checkpoint_path, best_model_path)\n        \ndef load_ckp(checkpoint_fpath, model, optimizer):\n    \"\"\"\n    checkpoint_path: path to save checkpoint\n    model: model that we want to load checkpoint parameters into       \n    optimizer: optimizer we defined in previous training\n    \"\"\"\n    # load check point\n    checkpoint = torch.load(checkpoint_fpath)\n    # initialize state_dict from checkpoint to model\n    model.load_state_dict(checkpoint['state_dict'])\n#     # initialize optimizer from checkpoint to optimizer\n#     optimizer.load_state_dict(checkpoint['optimizer'])\n#     # initialize valid_loss_min from checkpoint to valid_loss_min\n#     valid_loss_min = checkpoint['valid_loss_min']\n#     # return model, optimizer, epoch value, min validation loss \n    return model #, optimizer, checkpoint['epoch'], valid_loss_min.item()","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:21:27.736797Z","iopub.execute_input":"2021-06-03T12:21:27.737143Z","iopub.status.idle":"2021-06-03T12:21:27.745189Z","shell.execute_reply.started":"2021-06-03T12:21:27.737112Z","shell.execute_reply":"2021-06-03T12:21:27.743962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1. Prepare Dataloader","metadata":{}},{"cell_type":"code","source":"'''\noutput of this dataset should satisfy the requirements of torchvision faster-rcnn model.\nmore details please refer to https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html\n'''\nclass GolfDataset(Dataset):\n    def __init__(self,image_dir, label_dir=None,phase='train'):\n        super(GolfDataset, self).__init__()\n        self.image_dir = image_dir\n        self.label_dir = label_dir\n        self.phase = phase\n        self.image_files = self._search_samples(image_dir)\n        self.transforms = get_transforms(phase)\n        \n    def __len__(self):\n        return len(self.image_files)\n    \n    def __getitem__(self,idx):\n#         pdb.set_trace()\n        image_file = self.image_files[idx]\n        sample_name = image_file.split('/')[-1].split('.')[0].split('_')[-1]\n        label_file = os.path.join(self.label_dir, 'label_'+ sample_name +'.txt')\n        \n        if not os.path.exists(image_file) or not os.path.exists(label_file):\n            print('Warnning, image file or lable file not exist, please check your data')\n            exit(-1)\n        \n#         pdb.set_trace()\n        image = Image.open(image_file)\n        image_arr = cv2.imread(image_file)\n        image_arr = cv2.cvtColor(image_arr, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image_arr /= 255.0\n        \n        if self.phase == 'train':\n            target = self._parse_label_info(image_arr, image_file, label_file, idx)\n        \n            if self.transforms:\n                sample = {\n                    'image': image_arr,\n                    'bboxes': target['boxes'],\n                    'labels': target['labels']\n                }\n                sample = self.transforms(**sample)\n                image = sample['image']\n\n            target['boxes'] = torch.stack(tuple(map(torch.tensor, \n                                                zip(*sample['bboxes'])))).permute(1, 0)\n        \n        if self.phase != 'train':\n            return torch.as_tensor(image_arr), image_file\n        else:\n            return image, target, image_file\n        \n    def _search_samples(self, image_dir):\n        sample_files= []\n        label_files = []\n        filenames = os.listdir(image_dir)\n        for filename in filenames:\n            sample_files.append(os.path.join(image_dir, filename))\n       \n        return sample_files\n    \n    def _parse_label_info(self, image, image_file, label_file, idx):\n        '''\n        lable format: \n        class_id, center_x, center_y, bbox_weight, bbox_height\n        pixel values are normalized\n        \n        more details please refers to yolov5 - https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data\n        '''\n        h, w, _ = image.shape\n        label_infos = open(label_file).readlines()\n        \n        boxes = []\n        class_ids = []\n        \n        for label_info in label_infos:\n            class_id, center_x, center_y, width, height = label_info.strip().split()\n            \n            # convert all infos into int value\n            class_id = int(class_id) + 1\n            center_x, center_y = int(w * float(center_x)), int(h * float(center_y))\n            width, height = int(w * float(width)), int(h * float(height))\n            \n            # boxes\n            top_left = (int(center_x - width/2), int(center_y - height / 2))\n            bottom_right = (int(center_x + width/2), int(center_y + height / 2))\n            \n            # delete wrong labels\n            if top_left[0] == bottom_right[0] or top_left[1] == bottom_right[1]:continue\n            box = [top_left[0], top_left[1], bottom_right[0], bottom_right[1]]\n                        \n            class_ids.append(class_id)\n            boxes.append(box)\n        \n        class_ids = torch.as_tensor(class_ids, dtype=torch.int64)\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        \n        # area\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n\n        # iscrowd\n        iscrowd = torch.zeros((len(class_ids),), dtype=torch.int64)\n        \n        # labels is a list with dict element\n        labels = {}\n        labels[\"boxes\"] = boxes\n        labels[\"labels\"] = class_ids\n        labels[\"area\"] = area\n        labels[\"iscrowd\"] = iscrowd\n        labels[\"image_id\"] = torch.tensor([idx])\n        \n        return labels","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:21:30.899785Z","iopub.execute_input":"2021-06-03T12:21:30.900152Z","iopub.status.idle":"2021-06-03T12:21:30.922941Z","shell.execute_reply.started":"2021-06-03T12:21:30.900121Z","shell.execute_reply":"2021-06-03T12:21:30.922138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''\nPrepare Dataloader\n'''\ntrain_image_dir = '/kaggle/input/golfdetection/train_images/train_images'\ntrain_label_dir = '/kaggle/input/golfdetection/train_labels/train_labels'\ntest_image_dir = '/kaggle/input/golfdetection/test_images/test_images'\n\ntrain_dataset = GolfDataset(train_image_dir, train_label_dir, phase='train')\ntest_dataset = GolfDataset(test_image_dir, None, phase='test')\n\ntrain_dataloader = DataLoader(train_dataset,batch_size=8,shuffle=True,num_workers=6, collate_fn=collate_fn) # num_workers set to 0 for debugging\ntest_dataloader = DataLoader(test_dataset,batch_size=1,shuffle=False,num_workers=1, collate_fn=collate_fn) # num_workers set to 0 for debugging","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:21:34.929530Z","iopub.execute_input":"2021-06-03T12:21:34.929941Z","iopub.status.idle":"2021-06-03T12:21:35.832310Z","shell.execute_reply.started":"2021-06-03T12:21:34.929904Z","shell.execute_reply":"2021-06-03T12:21:35.831448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''\nun-comment dataset block area in dataset and following codes for visualization test\n'''\n# for images, targets, image_ids in tqdm(train_dataloader):\n#     break\n# # a = 1\n# # # # images, targets, ids = next(iter(train_dataloader))\n# # # pdb.set_trace()\n# # # a = 1","metadata":{"execution":{"iopub.status.busy":"2021-06-03T08:16:07.020282Z","iopub.execute_input":"2021-06-03T08:16:07.02063Z","iopub.status.idle":"2021-06-03T08:16:07.02642Z","shell.execute_reply.started":"2021-06-03T08:16:07.020596Z","shell.execute_reply":"2021-06-03T08:16:07.025488Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Prepare Model","metadata":{}},{"cell_type":"code","source":"'''\nRe-define Model\n'''\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\nnum_classes = 4  # 3 class + background\n\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:21:40.272450Z","iopub.execute_input":"2021-06-03T12:21:40.272787Z","iopub.status.idle":"2021-06-03T12:21:49.339436Z","shell.execute_reply.started":"2021-06-03T12:21:40.272756Z","shell.execute_reply":"2021-06-03T12:21:49.338580Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Prepare Optimizers","metadata":{}},{"cell_type":"code","source":"model.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=1, gamma=0.2)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:21:51.281704Z","iopub.execute_input":"2021-06-03T12:21:51.282044Z","iopub.status.idle":"2021-06-03T12:21:55.573777Z","shell.execute_reply.started":"2021-06-03T12:21:51.282000Z","shell.execute_reply":"2021-06-03T12:21:55.572975Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Training","metadata":{}},{"cell_type":"code","source":"num_epochs = 1\ntrain_loss_min = 0.9\ntotal_train_loss = []\n\n\ncheckpoint_dir = '/kaggle/working/checkpoint_rcnn'\nbest_model_path = '/kaggle/working/checkpoint_rcnn/best_rcnn.pt'\nif not os.path.exists(checkpoint_dir):\n    os.makedirs(checkpoint_dir)\n\n# tranining process\nfor epoch in range(num_epochs):\n    print(f'Epoch :{epoch + 1}')\n    train_loss = []\n    model.train()\n    iter = 0\n    \n    # iters for each epoch\n    for images, targets, image_ids in tqdm(train_dataloader):\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        losses = sum(loss for loss in loss_dict.values())\n        train_loss.append(losses.item())\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n        \n        if iter % 500 == 0:\n            print(losses.item())\n        if iter == 800:   # since this is a demo, we want to save training time\n            break\n        iter += 1\n    \n    # update learning rate\n    lr_scheduler.step()\n    \n    # log and model saveing\n    epoch_train_loss = np.mean(train_loss)\n    total_train_loss.append(epoch_train_loss)\n    print(f'Epoch train loss is {epoch_train_loss}')\n    \n    checkpoint = {\n            'epoch': epoch + 1,\n            'train_loss_min': epoch_train_loss,\n            'state_dict': model.state_dict(),\n            'optimizer': optimizer.state_dict(),\n        }\n    \n    checkpoint_path = os.path.join(checkpoint_dir, str(epoch)+'_rcnn.pt')\n    save_ckp(checkpoint, False, checkpoint_path, best_model_path)\n    ## TODO: save the model if validation loss has decreased\n    if epoch_train_loss <= train_loss_min:\n            print('Train loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(train_loss_min,epoch_train_loss))\n            # save checkpoint as best model\n            save_ckp(checkpoint, True, checkpoint_path, best_model_path)\n            train_loss_min = epoch_train_loss","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-06-03T12:22:04.763591Z","iopub.execute_input":"2021-06-03T12:22:04.763898Z","iopub.status.idle":"2021-06-03T12:42:34.123465Z","shell.execute_reply.started":"2021-06-03T12:22:04.763870Z","shell.execute_reply":"2021-06-03T12:42:34.122475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.title('Train Loss')\nplt.plot(total_train_loss)\nplt.xlabel('epochs')\nplt.ylabel('loss')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:43:17.068983Z","iopub.execute_input":"2021-06-03T12:43:17.069342Z","iopub.status.idle":"2021-06-03T12:43:17.073158Z","shell.execute_reply.started":"2021-06-03T12:43:17.069310Z","shell.execute_reply":"2021-06-03T12:43:17.071880Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5. Evaluation","metadata":{}},{"cell_type":"code","source":"'''\nSince Only train 800 iters (less than one epoch)\nSo we use training data for evaluation, and the result will not perform well\n'''\nfor images, targets, image_files in tqdm(train_dataloader):\n    images = list(image.to(device) for image in images)\n    \n    model1 = load_ckp('/kaggle/working/checkpoint_rcnn/0_rcnn.pt', model, optimizer)\n    model.eval()\n    predictions = model(images)\n    sample_prediction = predictions[0]\n    \n    # score threshold 0.5\n    scores = sample_prediction['scores'].detach().cpu().numpy().tolist()\n    boxes = sample_prediction['boxes'].detach().cpu().numpy().tolist()\n    labels = sample_prediction['labels'].detach().cpu().numpy().tolist()\n    \n    image = cv2.imread(image_files[0])\n    image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)\n    \n    results = []\n    for idx, score in enumerate(scores):\n        if score < 0.5: continue\n        box = boxes[idx]\n        label = labels[idx]\n        \n        if label == 1:\n            color = (0,255,0)\n        if label == 2:\n            color = (255,255,0)\n        if label == 3:\n            color = (0,255,255)\n        \n        cv2.rectangle(image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 2)\n        \n        obj_result = {'label': label, 'box': box, 'score': score}\n        results.append(obj_result)\n    \n    plt.figure(figsize=(20,20))\n    plt.imshow(image)\n    \n    np.save('demo_result.npy', results)\n    print(results)\n    break\n","metadata":{"execution":{"iopub.status.busy":"2021-06-03T12:55:35.248756Z","iopub.execute_input":"2021-06-03T12:55:35.249120Z","iopub.status.idle":"2021-06-03T12:55:40.739548Z","shell.execute_reply.started":"2021-06-03T12:55:35.249082Z","shell.execute_reply":"2021-06-03T12:55:40.733976Z"},"trusted":true},"execution_count":null,"outputs":[]}]}