{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\n\nfrom PIL import Image\nfrom enum import Enum\nfrom pathlib import Path\n\nfrom dataclasses import dataclass\n\nimport torch\nimport torchvision\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nfrom albumentations import *\nfrom albumentations.pytorch import ToTensorV2\nimport cv2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def fasterrcnn_resnet101_fpn(pretrained=False, progress=True,\n                            num_classes=91, pretrained_backbone=False,\n                             trainable_backbone_layers=3, **kwargs):\n    assert trainable_backbone_layers <= 5 and trainable_backbone_layers >= 0\n    # dont freeze any layers if pretrained model or backbone is not used\n    if not (pretrained or pretrained_backbone):\n        trainable_backbone_layers = 5\n    if pretrained:\n        # no need to download the backbone if pretrained is set\n        pretrained_backbone = False\n    backbone = resnet_fpn_backbone('resnet152', pretrained_backbone)\n    model = FasterRCNN(backbone, num_classes, **kwargs)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"def initialize_model():\n    model = fasterrcnn_resnet101_fpn(pretrained=False)\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = initialize_model()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"save_model_path = \"/kaggle/input/res152best-model-epoch-10\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model.load_state_dict(torch.load(os.path.join(save_model_path, f\"best_model_epoch_{20}.pth\")))\nmodel.load_state_dict(torch.load(os.path.join(save_model_path, f\"best_model_epoch_9.pth\")))\nmodel.to(device)\nmodel.eval()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@dataclass\nclass DatasetArguments:\n    data_dir: Path\n    images_lists_dict: dict\n    labels_csv_file_name: str\n\n@dataclass\nclass DataLoaderArguments:\n    batch_size: int\n    num_workers: int\n    dataset_arguments: DatasetArguments","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"phase=\"test\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"root_data_dir = Path(\"/kaggle/input/global-wheat-detection/\")\nunlabeled_generated_images_path = Path(f\"/kaggle/input/global-wheat-detection/{phase}/\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_images_file_names(directory):\n    _, _, files = next(os.walk(directory))\n    return files\ntest_file_names = get_images_file_names(unlabeled_generated_images_path)\ntest_file_names = [x.split(\".\")[0] for x in test_file_names]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images_lists_dict = {\n    \"test\": test_file_names\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"prediction_dataset_arguments = DatasetArguments(\n    data_dir=root_data_dir,\n    images_lists_dict=images_lists_dict,\n    labels_csv_file_name=\"sample_submission.csv\",\n)\npredict_dataloaders_arguments = DataLoaderArguments(\n    batch_size=4,\n    num_workers=0,\n    dataset_arguments=prediction_dataset_arguments\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def transform_set():\n    transforms_dict = {\n        'test': get_test_transforms()\n    }\n    return transforms_dict\n\n\ndef get_test_transforms():\n    return Compose(\n        [\n            ToTensorV2(p=1.0),\n        ]\n    )\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ObjectDetectionDataset(Dataset):\n    def __init__(self, images_root_directory, images_list, labels_csv_file_name, phase, transforms):\n        super(ObjectDetectionDataset).__init__()\n        self.images_root_directory = images_root_directory\n        self.phase = phase\n        self.transforms = transforms\n        self.images_list = images_list\n        if self.phase in [\"train\", \"val\"]:\n            self.labels_dataframe = pd.read_csv(os.path.join(images_root_directory, labels_csv_file_name))\n\n    def __getitem__(self, item):\n        sample = {\n            \"local_image_id\": None,\n            \"image_id\": None,\n            \"labels\": None,\n            \"boxes\": None,\n            \"area\": None,\n            \"iscrowd\": None\n        }\n\n        image_id = self.images_list[item]\n        image_path = os.path.join(self.images_root_directory,\n                                  \"train\" if self.phase in [\"train\", \"val\"] else \"test\",\n                                  image_id + \".jpg\")\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        sample[\"local_image_id\"] = image_id\n        sample[\"image_id\"] = torch.tensor([item])\n        if self.phase in [\"train\", \"val\"]:\n            boxes = self.labels_dataframe[self.labels_dataframe.image_id == image_id].bbox.values.tolist()\n            boxes = [eval(box_i) for box_i in boxes]\n            areas = _areas(boxes)\n            boxes = _adjust_boxes_format(boxes)\n\n            sample[\"labels\"] = torch.ones((len(boxes),), dtype=torch.int64)\n            sample[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.float32)\n            sample[\"area\"] = torch.as_tensor(areas, dtype=torch.float32)\n            sample[\"iscrowd\"] = torch.zeros((len(boxes),), dtype=torch.int64)\n        if self.transforms is not None:\n            sample[\"image\"] = image\n            if self.phase in [\"train\", \"val\"]:\n                transformed_sample = self.transforms(image=sample[\"image\"],\n                                                     bboxes=sample[\"boxes\"],\n                                                     labels=sample[\"labels\"])\n                sample[\"boxes\"] = torch.as_tensor(transformed_sample[\"bboxes\"], dtype=torch.float32)\n            else:\n                transformed_sample = self.transforms(image=sample[\"image\"])\n            image = transformed_sample[\"image\"]\n            del sample[\"image\"]\n        return image, sample\n    def __len__(self):\n        return len(self.images_list)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_dataset(arguments):\n    dataset = ObjectDetectionDataset(arguments.data_dir,\n                                     arguments.images_lists_dict[arguments.phase],\n                                     arguments.labels_csv_file_name,\n                                     arguments.phase,\n                                     arguments.transforms)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_prediction_dataloader(arguments, input_size):\n    data_transforms = transform_set()\n    batch_size = arguments.batch_size\n    num_workers = arguments.num_workers\n    arguments.dataset_arguments.phase = phase\n    arguments.dataset_arguments.transforms = data_transforms[\"test\"]\n    image_datasets = create_dataset(arguments.dataset_arguments)\n    dataloader = DataLoader(image_datasets, batch_size=batch_size,\n                            shuffle=False,\n                            pin_memory=True,\n                            num_workers=num_workers,\n                            collate_fn=collate_fn)\n    return dataloader\n\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader = create_prediction_dataloader(predict_dataloaders_arguments, None)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"detection_threshold=0.45\ndef format_prediction_string(boxes, scores):\n    pred_strings = []\n    for j in zip(scores, boxes):\n        pred_strings.append(\"{0:.4f} {1} {2} {3} {4}\".format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))\n\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results = []","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for images, sample in dataloader:\n    image_ids = [x[\"local_image_id\"] for x in sample]\n    images = list(image.to(device) for image in images)\n    outputs = model(images)\n\n    for i, image in enumerate(images):\n\n        boxes = outputs[i]['boxes'].data.cpu().numpy()\n        scores = outputs[i]['scores'].data.cpu().numpy()\n        \n        boxes = boxes[scores >= detection_threshold].astype(np.int32)\n        scores = scores[scores >= detection_threshold]\n        image_id = image_ids[i]\n        \n        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n        \n        result = {\n            'image_id': image_id,\n            'PredictionString': format_prediction_string(boxes, scores) if boxes.shape[0] > 0 else \"\"\n        }\n\n        \n        results.append(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.head()","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}