{"cells":[{"metadata":{},"cell_type":"markdown","source":"# DataSet\n\n## Define your own transform method","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import random\nfrom torchvision.transforms import functional as F\n\n\nclass Compose(object):\n def __init__(self, transforms):\n self.transforms = transforms\n\n def __call__(self, image, target):\n for t in self.transforms: \n image, target = t(image, target)\n return image, target\n\n\nclass ToTensor(object):\n def __call__(self, image, target):\n image = F.to_tensor(image)\n return image, target\n\n\nclass RandomHorizontalFlip(object):\n def __init__(self, prob=0.5):\n self.prob = prob\n\n def __call__(self, image, target):\n if random.random() < self.prob:\n height, width = image.shape[-2:]\n image = image.flip(-1)\n bbox = target[\"boxes\"]\n\n bbox[:, [0, 2]] = width - bbox[:, [2, 0]]\n target[\"boxes\"] = bbox\n return image, target","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visual bounding box","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import collections\nimport PIL.ImageDraw as ImageDraw\nimport PIL.ImageFont as ImageFont\nimport numpy as np\n\nSTANDARD_COLORS = [\n 'Pink', 'Green', 'SandyBrown',\n 'SeaGreen', 'Silver', 'SkyBlue', 'White',\n 'WhiteSmoke', 'YellowGreen'\n]\n\n\ndef filter_low_thresh(boxes, scores, classes, category_index, thresh, box_to_display_str_map, box_to_color_map, col):\n for i in range(boxes.shape[0]):\n if scores[i] > thresh:\n box = tuple(boxes[i].tolist()) # numpy -> list -> tuple\n if classes[i] in category_index.keys():\n class_name = category_index[classes[i]]\n else:\n class_name = 'N/A'\n display_str = str(class_name)\n display_str = '{}: {}%'.format(display_str, int(100 * scores[i]))\n box_to_display_str_map[box].append(display_str)\n box_to_color_map[box] = STANDARD_COLORS[col]\n else:\n break # Scores have been sorted\n\n\ndef draw_text(draw, box_to_display_str_map, box, left, right, top, bottom, color):\n try:\n font = ImageFont.truetype('arial.ttf', 24)\n except IOError:\n font = ImageFont.load_default()\n\n display_str_heights = [font.getsize(ds)[1] for ds in box_to_display_str_map[box]]\n total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)\n\n if top > total_display_str_height:\n text_bottom = top\n else:\n text_bottom = bottom + total_display_str_height\n for display_str in box_to_display_str_map[box][::-1]:\n text_width, text_height = font.getsize(display_str)\n margin = np.ceil(0.05 * text_height)\n draw.rectangle([(left, text_bottom - text_height - 2 * margin),\n (left + text_width, text_bottom)], fill=color)\n draw.text((left + margin, text_bottom - text_height - margin),\n display_str,\n fill='black',\n font=font)\n text_bottom -= text_height - 2 * margin\n\n\ndef draw_box(image, boxes, classes, scores, category_index, thresh=0.5, line_thickness=8):\n box_to_display_str_map = collections.defaultdict(list)\n box_to_color_map = collections.defaultdict(str)\n \n col = int(random.random() * len(STANDARD_COLORS))\n filter_low_thresh(boxes, scores, classes, category_index, thresh, box_to_display_str_map, box_to_color_map, col)\n\n draw = ImageDraw.Draw(image)\n im_width, im_height = image.size\n for box, color in box_to_color_map.items():\n xmin, ymin, xmax, ymax = box\n (left, right, top, bottom) = (xmin * 1, xmax * 1,\n ymin * 1, ymax * 1)\n draw.line([(left, top), (left, bottom), (right, bottom),\n (right, top), (left, top)], width=line_thickness, fill=color)\n draw_text(draw, box_to_display_str_map, box, left, right, top, bottom, color)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\n## WheatDataSet","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.utils.data import Dataset\nimport os\nimport pandas as pd\nimport json\n\n\nclass WheatDataSet(Dataset):\n def __init__(self, root, transforms):\n self.root = os.path.join(root)\n\n self.img_root = os.path.join(self.root, 'train')\n self.csv_root = os.path.join(self.root, 'train.csv')\n\n data = pd.read_csv(self.csv_root)\n data = data.drop(['width', 'height', 'source'], axis=1)\n img_list = data['image_id'].tolist()\n object = data['bbox'].tolist()\n\n self.class_dict = {\"wheat\": 1}\n\n self.transforms = transforms\n img_set = []\n name_set = []\n if object is not None:\n for img, data in zip(img_list, object):\n if img not in name_set:\n name_set.append(img)\n img_set.append({'name': img, 'bbox': [data]})\n else:\n idx = name_set.index(img)\n img_set[idx]['bbox'].append(data)\n else:\n for img in img_list:\n if img not in name_set:\n name_set.append(img)\n img_set.append({'name': img})\n self.img_set = img_set\n\n def __len__(self):\n return len(self.img_set)\n\n def __getitem__(self, idx):\n img_path = os.path.join(self.img_root, self.img_set[idx]['name'] + '.jpg')\n image = Image.open(img_path)\n\n obj = [i[1:-1].split(', ') for i in self.img_set[idx]['bbox']]\n temp = []\n for bbox in obj:\n t = [float(i) for i in bbox]\n temp.append(t)\n obj = temp\n\n boxes = []\n for o in obj:\n xmin = o[0]\n ymin = o[1]\n xmax = xmin + o[2]\n ymax = ymin + o[3]\n boxes.append([xmin, ymin, xmax, ymax])\n\n boxes = torch.as_tensor(boxes, dtype=torch.float32)\n area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n labels = torch.ones(boxes.shape[0], dtype=torch.int64)\n iscrowd = torch.zeros(boxes.shape[0], dtype=torch.int64)\n image_id = torch.tensor([idx])\n\n target = {\"boxes\": boxes, \"labels\": labels, \"image_id\": image_id, \"area\": area, \"iscrowd\": iscrowd}\n\n if self.transforms is not None:\n image, target = self.transforms(image, target)\n \n image /= 255.0\n\n return image, target\n\n\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nimport torchvision.transforms as ts\nimport random\n\ndata_transform = {\n \"train\": Compose([ToTensor(),\n RandomHorizontalFlip(0.5)]),\n \"val\": Compose([ToTensor()])\n}\n\n# load train data set\ntrain_data_set = WheatDataSet('../input/global-wheat-detection', data_transform[\"train\"])\nfor index in random.sample(range(0, len(train_data_set)), k=1):\n img, target = train_data_set[index]\n img = ts.ToPILImage()(img)\n draw_box(img,\n target[\"boxes\"].numpy(),\n target[\"labels\"].numpy(),\n [1 for i in range(len(target[\"labels\"].numpy()))],\n {1: 'Wheat'},\n thresh=0.5,\n line_thickness=5)\n plt.imshow(img)\n plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model\n## Backbone\nCreate a ResNet50 FPN Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\nimport torch\nfrom torch import Tensor\nfrom collections import OrderedDict\nimport torch.nn.functional as F\n\nfrom torch.jit.annotations import List, Dict\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, in_channel, out_channel, stride=1, downsample=None, norm_layer=None):\n super(Bottleneck, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n\n self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,\n kernel_size=1, stride=1, bias=False) # squeeze channels\n self.bn1 = norm_layer(out_channel)\n \n self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,\n kernel_size=3, stride=stride, bias=False, padding=1)\n self.bn2 = norm_layer(out_channel)\n \n self.conv3 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel * self.expansion,\n kernel_size=1, stride=1, bias=False) # unsqueeze channels\n self.bn3 = norm_layer(out_channel * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n\n def forward(self, x):\n identity = x\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n\n def __init__(self, block, blocks_num, num_classes=1000, include_top=True, norm_layer=None):\n super(ResNet, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n self._norm_layer = norm_layer\n\n self.include_top = include_top\n self.in_channel = 64\n\n self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,\n padding=3, bias=False)\n self.bn1 = norm_layer(self.in_channel)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, blocks_num[0])\n self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)\n self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)\n self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)\n if self.include_top:\n self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # output size = (1, 1)\n self.fc = nn.Linear(512 * block.expansion, num_classes)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n\n def _make_layer(self, block, channel, block_num, stride=1):\n norm_layer = self._norm_layer\n downsample = None\n if stride != 1 or self.in_channel != channel * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),\n norm_layer(channel * block.expansion))\n\n layers = []\n layers.append(block(self.in_channel, channel, downsample=downsample,\n stride=stride, norm_layer=norm_layer))\n self.in_channel = channel * block.expansion\n\n for _ in range(1, block_num):\n layers.append(block(self.in_channel, channel, norm_layer=norm_layer))\n\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n\n if self.include_top:\n x = self.avgpool(x)\n x = torch.flatten(x, 1)\n x = self.fc(x)\n\n return x\n\nclass IntermediateLayerGetter(nn.ModuleDict):\n __annotations__ = {\n \"return_layers\": Dict[str, str],\n }\n\n def __init__(self, model, return_layers):\n if not set(return_layers).issubset([name for name, _ in model.named_children()]):\n raise ValueError(\"return_layers are not present in model\")\n\n orig_return_layers = return_layers\n return_layers = {k: v for k, v in return_layers.items()}\n layers = OrderedDict()\n\n for name, module in model.named_children():\n layers[name] = module\n if name in return_layers:\n del return_layers[name]\n if not return_layers:\n break\n\n super(IntermediateLayerGetter, self).__init__(layers)\n self.return_layers = orig_return_layers\n\n def forward(self, x):\n out = OrderedDict()\n \n for name, module in self.named_children():\n x = module(x)\n if name in self.return_layers:\n out_name = self.return_layers[name]\n out[out_name] = x\n return out\n\n\nclass FeaturePyramidNetwork(nn.Module):\n \n\n def __init__(self, in_channels_list, out_channels, extra_blocks=None):\n super(FeaturePyramidNetwork, self).__init__()\n \n self.inner_blocks = nn.ModuleList()\n \n self.layer_blocks = nn.ModuleList()\n for in_channels in in_channels_list:\n if in_channels == 0:\n continue\n inner_block_module = nn.Conv2d(in_channels, out_channels, 1)\n layer_block_module = nn.Conv2d(out_channels, out_channels, 3, padding=1)\n self.inner_blocks.append(inner_block_module)\n self.layer_blocks.append(layer_block_module)\n\n for m in self.children():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_uniform_(m.weight, a=1)\n nn.init.constant_(m.bias, 0)\n\n self.extra_blocks = extra_blocks\n\n def get_result_from_inner_blocks(self, x, idx):\n # type: (Tensor, int)\n \n num_blocks = 0\n for m in self.inner_blocks:\n num_blocks += 1\n if idx < 0:\n idx += num_blocks\n i = 0\n out = x\n for module in self.inner_blocks:\n if i == idx:\n out = module(x)\n i += 1\n return out\n\n def get_result_from_layer_blocks(self, x, idx):\n # type: (Tensor, int)\n \n num_blocks = 0\n for m in self.layer_blocks:\n num_blocks += 1\n if idx < 0:\n idx += num_blocks\n i = 0\n out = x\n for module in self.layer_blocks:\n if i == idx:\n out = module(x)\n i += 1\n return out\n\n def forward(self, x):\n # type: (Dict[str, Tensor])\n\n names = list(x.keys())\n x = list(x.values())\n\n last_inner = self.inner_blocks[-1](x[-1])\n\n results = []\n \n results.append(self.layer_blocks[-1](last_inner))\n\n for idx in range(len(x) - 2, -1, -1):\n inner_lateral = self.get_result_from_inner_blocks(x[idx], idx)\n feat_shape = inner_lateral.shape[-2:]\n inner_top_down = F.interpolate(last_inner, size=feat_shape, mode=\"nearest\")\n last_inner = inner_lateral + inner_top_down\n results.insert(0, self.get_result_from_layer_blocks(last_inner, idx))\n\n if self.extra_blocks is not None:\n results, names = self.extra_blocks(results, names)\n\n # make it back an OrderedDict\n out = OrderedDict([(k, v) for k, v in zip(names, results)])\n\n return out\n\n\nclass LastLevelMaxPool(torch.nn.Module):\n\n def forward(self, x, names):\n # type: (List[Tensor], List[str])\n names.append(\"pool\")\n x.append(F.max_pool2d(x[-1], 1, 2, 0))\n return x, names\n\n\nclass BackboneWithFPN(nn.Module):\n\n def __init__(self, backbone, return_layers, in_channels_list, out_channels):\n super(BackboneWithFPN, self).__init__()\n self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)\n self.fpn = FeaturePyramidNetwork(\n in_channels_list=in_channels_list,\n out_channels=out_channels,\n extra_blocks=LastLevelMaxPool(),\n )\n self.out_channels = out_channels\n\n def forward(self, x):\n x = self.body(x)\n x = self.fpn(x)\n return x\n\n\ndef resnet50_fpn_backbone():\n resnet_backbone = ResNet(Bottleneck, [3, 4, 6, 3],\n include_top=False)\n for name, parameter in resnet_backbone.named_parameters():\n if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:\n parameter.requires_grad_(False)\n\n return_layers = {'layer1': '0', 'layer2': '1', 'layer3': '2', 'layer4': '3'}\n\n in_channels_stage2 = resnet_backbone.in_channel // 8\n in_channels_list = [\n in_channels_stage2, # layer1 out_channel=256\n in_channels_stage2 * 2, # layer2 out_channel=512\n in_channels_stage2 * 4, # layer3 out_channel=1024\n in_channels_stage2 * 8, # layer4 out_channel=2048\n ]\n out_channels = 256\n return BackboneWithFPN(resnet_backbone, return_layers, in_channels_list, out_channels)\n\n\nprint(resnet50_fpn_backbone())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Bounding box","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.jit.annotations import Tuple\nfrom torch import Tensor\nimport torch._ops\nimport torchvision\n\n\ndef nms(boxes, scores, iou_threshold):\n # type: (Tensor, Tensor, float)\n return torch.ops.torchvision.nms(boxes, scores, iou_threshold)\n\n\ndef batched_nms(boxes, scores, idxs, iou_threshold):\n # type: (Tensor, Tensor, Tensor, float)\n if boxes.numel() == 0:\n return torch.empty((0,), dtype=torch.int64, device=boxes.device)\n\n max_coordinate = boxes.max()\n offsets = idxs.to(boxes) * (max_coordinate + 1)\n\n boxes_for_nms = boxes + offsets[:, None]\n keep = nms(boxes_for_nms, scores, iou_threshold)\n return keep\n\n\ndef remove_small_boxes(boxes, min_size):\n # type: (Tensor, float)\n ws, hs = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]\n keep = (ws >= min_size) & (hs >= min_size)\n keep = keep.nonzero().squeeze(1)\n return keep\n\n\ndef clip_boxes_to_image(boxes, size):\n # type: (Tensor, Tuple[int, int])\n\n dim = boxes.dim()\n boxes_x = boxes[..., 0::2]\n boxes_y = boxes[..., 1::2]\n height, width = size\n\n if torchvision._is_tracing():\n boxes_x = torch.max(boxes_x, torch.tensor(0, dtype=boxes.dtype, device=boxes.device))\n boxes_x = torch.min(boxes_x, torch.tensor(width, dtype=boxes.dtype, device=boxes.device))\n boxes_y = torch.max(boxes_y, torch.tensor(0, dtype=boxes.dtype, device=boxes.device))\n boxes_y = torch.min(boxes_y, torch.tensor(height, dtype=boxes.dtype, device=boxes.device))\n else:\n boxes_x = boxes_x.clamp(min=0, max=width)\n boxes_y = boxes_y.clamp(min=0, max=width)\n\n clipped_boxes = torch.stack((boxes_x, boxes_y), dim=dim)\n return clipped_boxes.reshape(boxes.shape)\n\n\ndef box_area(boxes):\n return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])\n\n\ndef box_iou(boxes1, boxes2):\n area1 = box_area(boxes1)\n area2 = box_area(boxes2)\n\n lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])\n rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])\n\n wh = (rb - lt).clamp(min=0)\n inter = wh[:, :, 0] * wh[:, :, 1]\n\n iou = inter / (area1[:, None] + area2 - inter)\n return iou","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport math\nfrom torch.jit.annotations import List, Tuple\nfrom torch import Tensor\n\n\ndef zeros_like(tensor, dtype):\n # type: (Tensor, int) -> Tensor\n return torch.zeros_like(tensor, dtype=dtype, layout=tensor.layout,\n device=tensor.device, pin_memory=tensor.is_pinned())\n\n\n# @torch.jit.script\nclass BalancedPositiveNegativeSampler(object):\n def __init__(self, batch_size_per_image, positive_fraction):\n # type: (int, float)\n self.batch_size_per_image = batch_size_per_image\n self.positive_fraction = positive_fraction\n\n def __call__(self, matched_idxs):\n # type: (List[Tensor])\n pos_idx = []\n neg_idx = []\n for matched_idxs_per_image in matched_idxs:\n positive = torch.nonzero(matched_idxs_per_image >= 1).squeeze(1)\n negative = torch.nonzero(matched_idxs_per_image == 0).squeeze(1)\n\n num_pos = int(self.batch_size_per_image * self.positive_fraction)\n # If the number of positive samples is not enough, all positive samples will be used directly\n num_pos = min(positive.numel(), num_pos)\n \n num_neg = self.batch_size_per_image - num_pos\n # Same as positive samples\n num_neg = min(negative.numel(), num_neg)\n\n # Randomly select a specified number of positive and negative samples\n perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]\n perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]\n\n pos_idx_per_image = positive[perm1]\n neg_idx_per_image = negative[perm2]\n\n pos_idx_per_image_mask = zeros_like(\n matched_idxs_per_image, dtype=torch.uint8\n )\n neg_idx_per_image_mask = zeros_like(\n matched_idxs_per_image, dtype=torch.uint8\n )\n\n pos_idx_per_image_mask[pos_idx_per_image] = torch.tensor(1, dtype=torch.uint8)\n neg_idx_per_image_mask[neg_idx_per_image] = torch.tensor(1, dtype=torch.uint8)\n\n pos_idx.append(pos_idx_per_image_mask)\n neg_idx.append(neg_idx_per_image_mask)\n return pos_idx, neg_idx\n\n\n# @torch.jit.script\ndef encode_boxes(reference_boxes, proposals, weights):\n # type: (torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor\n\n wx = weights[0]\n wy = weights[1]\n ww = weights[2]\n wh = weights[3]\n\n proposals_x1 = proposals[:, 0].unsqueeze(1)\n proposals_y1 = proposals[:, 1].unsqueeze(1)\n proposals_x2 = proposals[:, 2].unsqueeze(1)\n proposals_y2 = proposals[:, 3].unsqueeze(1)\n\n reference_boxes_x1 = reference_boxes[:, 0].unsqueeze(1)\n reference_boxes_y1 = reference_boxes[:, 1].unsqueeze(1)\n reference_boxes_x2 = reference_boxes[:, 2].unsqueeze(1)\n reference_boxes_y2 = reference_boxes[:, 3].unsqueeze(1)\n\n ex_widths = proposals_x2 - proposals_x1\n ex_heights = proposals_y2 - proposals_y1\n\n ex_ctr_x = proposals_x1 + 0.5 * ex_widths\n ex_ctr_y = proposals_y1 + 0.5 * ex_heights\n\n gt_widths = reference_boxes_x2 - reference_boxes_x1\n gt_heights = reference_boxes_y2 - reference_boxes_y1\n gt_ctr_x = reference_boxes_x1 + 0.5 * gt_widths\n gt_ctr_y = reference_boxes_y1 + 0.5 * gt_heights\n\n targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths\n targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights\n targets_dw = ww * torch.log(gt_widths / ex_widths)\n targets_dh = wh * torch.log(gt_heights / ex_heights)\n\n targets = torch.cat((targets_dx, targets_dy, targets_dw, targets_dh), dim=1)\n return targets\n\n\n# @torch.jit.script\nclass BoxCoder(object):\n def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16)):\n # type: (Tuple[float, float, float, float], float)\n self.weights = weights\n self.bbox_xform_clip = bbox_xform_clip\n\n def encode(self, reference_boxes, proposals):\n # type: (List[Tensor], List[Tensor])\n\n \"\"\"\n Calculate the expression parameter\n\n \"\"\"\n\n # Count the number of positive and negative samples of each image, \n # so as to facilitate the later splicing together and separate after processing\n boxes_per_image = [len(b) for b in reference_boxes]\n reference_boxes = torch.cat(reference_boxes, dim=0)\n proposals = torch.cat(proposals, dim=0)\n\n targets = self.encode_single(reference_boxes, proposals)\n return targets.split(boxes_per_image, 0)\n\n def encode_single(self, reference_boxes, proposals):\n\n dtype = reference_boxes.dtype\n device = reference_boxes.device\n weights = torch.as_tensor(self.weights, dtype=dtype, device=device)\n targets = encode_boxes(reference_boxes, proposals, weights)\n\n return targets\n\n def decode(self, rel_codes, boxes):\n # type: (Tensor, List[Tensor])\n assert isinstance(boxes, (list, tuple))\n assert isinstance(rel_codes, torch.Tensor)\n boxes_per_image = [b.size(0) for b in boxes]\n concat_boxes = torch.cat(boxes, dim=0)\n\n box_sum = 0\n for val in boxes_per_image:\n box_sum += val\n pred_boxes = self.decode_single(\n rel_codes.reshape(box_sum, -1), concat_boxes\n )\n return pred_boxes.reshape(box_sum, -1, 4)\n\n def decode_single(self, rel_codes, boxes):\n boxes = boxes.to(rel_codes.dtype)\n\n # xmin, ymin, xmax, ymax\n widths = boxes[:, 2] - boxes[:, 0] # The widths of anchor\n heights = boxes[:, 3] - boxes[:, 1] # The heights of anchor\n ctr_x = boxes[:, 0] + 0.5 * widths # Anchor center X coordinate\n ctr_y = boxes[:, 1] + 0.5 * heights # Anchor center Y coordinate\n\n wx, wy, ww, wh = self.weights # The default is 1\n dx = rel_codes[:, 0::4] / wx # Regression parameters of the central coordinate X of anchors\n dy = rel_codes[:, 1::4] / wy # Regression parameters of the central coordinate Y of anchors\n dw = rel_codes[:, 2::4] / ww # Regression parameters for predicting the width of anchors\n dh = rel_codes[:, 3::4] / wh # Regression parameters for predicting the heights of anchors\n\n dw = torch.clamp(dw, max=self.bbox_xform_clip)\n dh = torch.clamp(dh, max=self.bbox_xform_clip)\n\n pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]\n pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]\n pred_w = torch.exp(dw) * widths[:, None]\n pred_h = torch.exp(dh) * heights[:, None]\n\n # xmin\n pred_boxes1 = pred_ctr_x - torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w\n # ymin\n pred_boxes2 = pred_ctr_y - torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_h\n # xmax\n pred_boxes3 = pred_ctr_x + torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w\n # ymax\n pred_boxes4 = pred_ctr_y + torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_h\n pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=2).flatten(1)\n return pred_boxes\n\n\n# @torch.jit.script\nclass Matcher(object):\n BELOW_LOW_THRESHOLD = -1\n BETWEEN_THRESHOLDS = -2\n\n __annotations__ = {\n 'BELOW_LOW_THRESHOLD': int,\n 'BETWEEN_THRESHOLDS': int,\n }\n\n def __init__(self, high_threshold, low_threshold, allow_low_quality_matches=False):\n # type: (float, float, bool)\n self.BELOW_LOW_THRESHOLD = -1\n self.BETWEEN_THRESHOLDS = -2\n assert low_threshold <= high_threshold\n self.high_threshold = high_threshold\n self.low_threshold = low_threshold\n self.allow_low_quality_matches = allow_low_quality_matches\n\n def __call__(self, match_quality_matrix):\n \"\"\"\n Calculate the maximum IOU value of anchors matching each gtboxes, and record the index\n \"\"\"\n if match_quality_matrix.numel() == 0:\n # empty targets or proposals not supported during training\n if match_quality_matrix.shape[0] == 0:\n raise ValueError(\n \"No ground-truth boxes available for one of the images \"\n \"during training\")\n else:\n raise ValueError(\n \"No proposal boxes available for one of the images \"\n \"during training\")\n matched_vals, matches = match_quality_matrix.max(dim=0) # the dimension to reduce.\n if self.allow_low_quality_matches:\n all_matches = matches.clone()\n else:\n all_matches = None\n\n below_low_threshold = matched_vals < self.low_threshold\n between_thresholds = (matched_vals >= self.low_threshold) & (\n matched_vals < self.high_threshold\n )\n \n matches[below_low_threshold] = torch.tensor(self.BELOW_LOW_THRESHOLD) # -1\n\n matches[between_thresholds] = torch.tensor(self.BETWEEN_THRESHOLDS) # -2\n\n if self.allow_low_quality_matches:\n assert all_matches is not None\n self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)\n\n return matches\n\n def set_low_quality_matches_(self, matches, all_matches, match_quality_matrix):\n # For each gt boxes, find the anchor with the largest IOU\n # highest_quality_foreach_gt is the maximum IOU value matched\n highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)\n # Find the largest anchor index of each gt box and its IOU.\n # The largest IOU matched by a gt may have multiple anchors\n gt_pred_pairs_of_highest_quality = torch.nonzero(\n match_quality_matrix == highest_quality_foreach_gt[:, None]\n )\n pre_inds_to_update = gt_pred_pairs_of_highest_quality[:, 1]\n # Keep the index of the anchor that matches the gt maximum IOU, even if IOU is lower than the set threshold\n matches[pre_inds_to_update] = all_matches[pre_inds_to_update]\n\n\ndef smooth_l1_loss(input, target, beta: float = 1./9, size_average: bool = True):\n n = torch.abs(input - target)\n cond = n < beta\n loss = torch.where(cond, 0.5 * n ** 2 / beta, n - 0.5 * beta)\n if size_average:\n return loss.mean()\n return loss.sum()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.jit.annotations import List, Tuple\nfrom torch import Tensor\n\n\n# @torch.jit.script\nclass ImageList(object):\n def __init__(self, tensors, image_sizes):\n # type: (Tensor, List[Tuple[int, int]])\n\n self.tensors = tensors\n self.image_sizes = image_sizes\n\n def to(self, device):\n # type: (Device)\n cast_tensor = self.tensors.to(device)\n return ImageList(cast_tensor, self.image_sizes)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## RCNNTransform ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn, Tensor\nimport math\nfrom torch.jit.annotations import List, Tuple, Dict, Optional\nimport torchvision\n\n\nclass GeneralizedRCNNTransform(nn.Module):\n def __init__(self, min_size, max_size, image_mean, image_std):\n super(GeneralizedRCNNTransform, self).__init__()\n if not isinstance(min_size, (list, tuple)):\n min_size = (min_size, )\n self.min_size = min_size\n self.max_size = max_size\n self.image_mean = image_mean\n self.image_std = image_std\n\n def normalize(self, image):\n dtype, device = image.dtype, image.device\n mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)\n std = torch.as_tensor(self.image_std, dtype=dtype, device=device)\n return (image - mean[:, None, None]) / std[:, None, None]\n\n def torch_choice(self, l):\n # type: (List[int])\n index = int(torch.empty(1).uniform_(0., float(len(l))).item())\n return l[index]\n\n def resize(self, image, target):\n # type: (Tensor, Optional[Dict[str, Tensor]])\n \"\"\"\n Zoom the picture to the specified size range, and zoom bboxes information correspondingly\n Args:\n image: Input image\n target: Enter information about the picture（Include bboxes information）\n \"\"\"\n # image shape is [channel, height, width]\n h, w = image.shape[-2:]\n im_shape = torch.tensor(image.shape[-2:])\n min_size = float(torch.min(im_shape)) # Gets the minimum value in height and width\n max_size = float(torch.max(im_shape)) # Gets the maximum value in height and width\n if self.training:\n size = float(self.torch_choice(self.min_size)) # Specifies the minimum side length of the input picture\n else:\n # FIXME assume for now that testing uses the largest scale\n size = float(self.min_size[-1])\n # Calculates the zoom ratio based on the specified minimum side length and the minimum edge length of the picture\n scale_factor = size / min_size \n\n if max_size * scale_factor > self.max_size:\n scale_factor = self.max_size / max_size\n\n # image -> image[None] [C, H, W] -> [1, C, H, W]\n image = nn.functional.interpolate(\n image[None], scale_factor=scale_factor, mode='bilinear', align_corners=False\n )[0]\n\n if target is None:\n return image, target\n\n bbox = target[\"boxes\"]\n bbox = resize_boxes(bbox, (h, w), image.shape[-2:])\n target[\"boxes\"] = bbox\n\n return image, target\n\n @torch.jit.unused\n def _onnx_batch_images(self, images, size_divisible=32):\n # type: (List[Tensor], int) -> Tensor\n max_size = []\n for i in range(images[0].dim()):\n max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64)\n max_size.append(max_size_i)\n stride = size_divisible\n max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64)\n max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64)\n max_size = tuple(max_size)\n\n padded_imgs = []\n for img in images:\n padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]\n padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))\n padded_imgs.append(padded_img)\n\n return torch.stack(padded_imgs)\n\n def max_by_axis(self, the_list):\n # type: (List[List[int]]) -> List[int]\n maxes = the_list[0]\n for sublist in the_list[1:]:\n for index, item in enumerate(sublist):\n maxes[index] = max(maxes[index], item)\n return maxes\n\n def batch_images(self, images, size_divisible=32):\n # type: (List[Tensor], int)\n \"\"\"\n Package a batch of images into a batch\n \"\"\"\n if torchvision._is_tracing():\n return self._onnx_batch_images(images, size_divisible)\n max_size = self.max_by_axis([list(img.shape) for img in images])\n\n stride = float(size_divisible)\n max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)\n max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)\n\n batch_shape = [len(images)] + max_size\n\n batched_imgs = images[0].new_full(batch_shape, 0)\n for img, pad_img in zip(images, batched_imgs):\n # copy_: Copies the elements from src into self tensor and returns self\n pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)\n return batched_imgs\n\n def postprocess(self, result, image_shapes, original_image_sizes):\n # type: (List[Dict[str, Tensor]], List[Tuple[int, int]], List[Tuple[int, int]])\n \"\"\"\n Post processing（It mainly restores bboxes to the original image scale）\n Args:\n result: list(dict), Prediction results of network, len(result) == batch_size\n image_shapes: list(torch.Size), Size after image preprocessing and scaling, len(image_shapes) == batch_size\n original_image_sizes: list(torch.Size), The original size of the image, len(original_image_sizes) == batch_size\n \"\"\"\n if self.training:\n return result\n for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):\n boxes = pred[\"boxes\"]\n boxes = resize_boxes(boxes, im_s, o_im_s) # Zoom bboxes back to the original image scale\n result[i][\"boxes\"] = boxes\n return result\n\n def __repr__(self):\n format_string = self.__class__.__name__ + '('\n _indent = '\\n '\n format_string += \"{0}Normalize(mean={1}, std={2})\".format(_indent, self.image_mean, self.image_std)\n format_string += \"{0}Resize(min_size={1}, max_size={2}, mode='bilinear')\".format(_indent, self.min_size,\n self.max_size)\n format_string += '\\n)'\n return format_string\n\n def forward(self, images, targets=None):\n # type: (List[Tensor], Optional[List[Dict[str, Tensor]]])\n images = [img for img in images]\n for i in range(len(images)):\n image = images[i]\n target_index = targets[i] if targets is not None else None\n\n if image.dim() != 3:\n raise ValueError(\"images is expected to be a list of 3d tensors \"\n \"of shape [C, H, W], got {}\".format(image.shape))\n image = self.normalize(image)\n image, target_index = self.resize(image, target_index)\n images[i] = image\n if targets is not None and target_index is not None:\n targets[i] = target_index\n # Record the image size after resizing\n image_sizes = [img.shape[-2:] for img in images]\n images = self.batch_images(images) # Package images into a batch\n image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], [])\n\n for image_size in image_sizes:\n assert len(image_size) == 2\n image_sizes_list.append((image_size[0], image_size[1]))\n\n image_list = ImageList(images, image_sizes_list)\n return image_list, targets\n\n\ndef resize_boxes(boxes, original_size, new_size):\n # type: (Tensor, List[int], List[int]) -> Tensor\n \"\"\"\n The boxes parameter is scaled according to the image scaling\n \"\"\"\n ratios = [\n torch.tensor(s, dtype=torch.float32, device=boxes.device) /\n torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)\n for s, s_orig in zip(new_size, original_size)\n ]\n ratios_height, ratios_width = ratios\n # Removes a tensor dimension, boxes [minibatch, 4]\n # Returns a tuple of all slices along a given dimension, already without it.\n xmin, ymin, xmax, ymax = boxes.unbind(1)\n xmin = xmin * ratios_width\n xmax = xmax * ratios_width\n ymin = ymin * ratios_height\n ymax = ymax * ratios_height\n return torch.stack((xmin, ymin, xmax, ymax), dim=1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## ROI","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn.functional as F\nfrom torch import Tensor\nfrom torch.jit.annotations import Optional, List, Dict, Tuple\n\n\ndef fastrcnn_loss(class_logits, box_regression, labels, regression_targets):\n # type: (Tensor, Tensor, List[Tensor], List[Tensor])\n \"\"\"\n Computes the loss for Faster R-CNN.\n \"\"\"\n labels = torch.cat(labels, dim=0)\n regression_targets = torch.cat(regression_targets, dim=0)\n\n # Calculate category loss information\n classification_loss = F.cross_entropy(class_logits, labels)\n\n sampled_pos_inds_subset = torch.nonzero(labels > 0).squeeze(1)\n\n labels_pos = labels[sampled_pos_inds_subset]\n\n # shape=[num_proposal, num_classes]\n N, num_classes = class_logits.shape\n box_regression = box_regression.reshape(N, -1, 4)\n\n # Calculate bounding box loss information\n box_loss = smooth_l1_loss(\n box_regression[sampled_pos_inds_subset, labels_pos],\n regression_targets[sampled_pos_inds_subset],\n beta=1 / 9,\n size_average=False,\n ) / labels.numel()\n\n return classification_loss, box_loss\n\n\nclass ROIHeads(torch.nn.Module):\n __annotations__ = {\n 'box_coder': BoxCoder,\n 'proposal_matcher': Matcher,\n 'fg_bg_sampler': BalancedPositiveNegativeSampler,\n }\n\n def __init__(self, box_roi_pool, box_head, box_predictor,\n # Faster R-CNN training\n fg_iou_thresh, bg_iou_thresh, batch_size_per_image, positive_fraction, bbox_reg_weights,\n # Faster R-CNN inference\n score_thresh, nms_thresh, detection_per_img):\n super(ROIHeads, self).__init__()\n\n self.box_similarity = box_iou\n self.proposal_matcher = Matcher(\n fg_iou_thresh,\n bg_iou_thresh,\n allow_low_quality_matches=False\n )\n self.fg_bg_sampler = BalancedPositiveNegativeSampler(\n batch_size_per_image,\n positive_fraction\n )\n if bbox_reg_weights is None:\n bbox_reg_weights = (10., 10., 5., 5.)\n self.box_coder = BoxCoder(bbox_reg_weights)\n\n self.box_roi_pool = box_roi_pool\n self.box_head = box_head\n self.box_predictor = box_predictor\n\n self.score_thresh = score_thresh\n self.nms_thresh = nms_thresh\n self.detection_per_img = detection_per_img\n\n def assign_targets_to_proposals(self, proposals, gt_boxes, gt_labels):\n # type: (List[Tensor], List[Tensor], List[Tensor])\n matched_idxs = []\n labels = []\n for proposals_in_image, gt_boxes_in_image, gt_labels_in_image in zip(proposals, gt_boxes, gt_labels):\n if gt_boxes_in_image.numel() == 0:\n device = proposals_in_image.device\n clamped_matched_idxs_in_image = torch.zeros(\n (proposals_in_image.shape[0],), dtype=torch.int64, device=device\n )\n labels_in_image = torch.zeros(\n (proposals_in_image.shape[0],), dtype=torch.int64, device=device\n )\n else:\n # Calculate IOU\n match_quality_matrix = box_iou(gt_boxes_in_image, proposals_in_image)\n matched_idxs_in_image = self.proposal_matcher(match_quality_matrix)\n\n clamped_matched_idxs_in_image = matched_idxs_in_image.clamp(min=0)\n labels_in_image = gt_labels_in_image[clamped_matched_idxs_in_image]\n labels_in_image = labels_in_image.to(dtype=torch.int64)\n\n # label background (below the low threshold)\n # background\n bg_inds = matched_idxs_in_image == self.proposal_matcher.BELOW_LOW_THRESHOLD\n labels_in_image[bg_inds] = torch.tensor(0)\n\n # Waste samples\n ignore_inds = matched_idxs_in_image == self.proposal_matcher.BETWEEN_THRESHOLDS # -2\n labels_in_image[ignore_inds] = torch.tensor(-1)\n\n matched_idxs.append(clamped_matched_idxs_in_image)\n labels.append(labels_in_image)\n return matched_idxs, labels\n\n def subsample(self, labels):\n # type: (List[Tensor])\n\n sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)\n sampled_inds = []\n for img_idx, (pos_inds_img, neg_inds_img) in enumerate(zip(sampled_pos_inds, sampled_neg_inds)):\n img_sampled_inds = torch.nonzero(pos_inds_img | neg_inds_img).squeeze(1)\n sampled_inds.append(img_sampled_inds)\n return sampled_inds\n\n def add_gt_proposals(self, proposals, gt_boxes):\n # type: (List[Tensor], List[Tensor])\n proposals = [\n torch.cat((proposal, gt_box))\n for proposal, gt_box in zip(proposals, gt_boxes)\n ]\n return proposals\n\n def DELTEME_all(self, the_list):\n # type: (List[bool])\n for i in the_list:\n if not i:\n return False\n return True\n\n def check_targets(self, targets):\n # type: (Optional[List[Dict[str, Tensor]]])\n assert targets is not None\n assert self.DELTEME_all([\"boxes\" in t for t in targets])\n assert self.DELTEME_all([\"labels\" in t for t in targets])\n\n def select_training_samples(self, proposals, targets):\n # type: (List[Tensor], Optional[List[Dict[str, Tensor]]])\n \n # Check whether the target is empty\n self.check_targets(targets)\n assert targets is not None\n dtype = proposals[0].dtype\n device = proposals[0].device\n\n gt_boxes = [t[\"boxes\"].to(dtype) for t in targets]\n gt_labels = [t[\"labels\"] for t in targets]\n\n # Splice gt_boxes to the back of the proposal\n proposals = self.add_gt_proposals(proposals, gt_boxes)\n\n matched_idxs, labels = self.assign_targets_to_proposals(proposals, gt_boxes, gt_labels)\n\n # Positive and negative samples are sampled according to the given number and proportion\n sampled_inds = self.subsample(labels)\n matched_gt_boxes = []\n num_images = len(proposals)\n\n for img_id in range(num_images):\n # Get the positive and negative sample index of each image\n img_sampled_inds = sampled_inds[img_id]\n proposals[img_id] = proposals[img_id][img_sampled_inds]\n labels[img_id] = labels[img_id][img_sampled_inds]\n matched_idxs[img_id] = matched_idxs[img_id][img_sampled_inds]\n\n gt_boxes_in_image = gt_boxes[img_id]\n if gt_boxes_in_image.numel() == 0:\n gt_boxes_in_image = torch.zeros((1, 4), dtype=dtype, device=device)\n matched_gt_boxes.append(gt_boxes_in_image[matched_idxs[img_id]])\n\n regression_targets = self.box_coder.encode(matched_gt_boxes, proposals)\n return proposals, matched_idxs, labels, regression_targets\n\n def postprocess_detections(self, class_logits, box_regression, proposals, image_shapes):\n # type: (Tensor, Tensor, List[Tensor], List[Tuple[int, int]])\n device = class_logits.device\n num_classes = class_logits.shape[-1]\n\n boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals]\n pred_boxes = self.box_coder.decode(box_regression, proposals)\n\n pred_scores = F.softmax(class_logits, -1)\n\n pred_boxes_list = pred_boxes.split(boxes_per_image, 0)\n pred_scores_list = pred_scores.split(boxes_per_image, 0)\n\n all_boxes = []\n all_scores = []\n all_labels = []\n for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes):\n # Cut the predicted boxes information and\n # adjust the out of bounds coordinates to the image boundary\n boxes = clip_boxes_to_image(boxes, image_shape)\n\n # create labels for each prediction\n labels = torch.arange(num_classes, device=device)\n labels = labels.view(1, -1).expand_as(scores)\n\n # Remove all information with index 0 (0 represents background)\n boxes = boxes[:, 1:]\n scores = scores[:, 1:]\n labels = labels[:, 1:]\n\n boxes = boxes.reshape(-1, 4)\n scores = scores.reshape(-1)\n labels = labels.reshape(-1)\n\n # Remove low probability targets\n inds = torch.nonzero(scores > self.score_thresh).squeeze(1)\n boxes, scores, labels = boxes[inds], scores[inds], labels[inds]\n\n # Remove small target\n keep = remove_small_boxes(boxes, min_size=1e-2)\n boxes, scores, labels = boxes[keep], scores[keep], labels[keep]\n\n # After NMS processing is executed, \n # the results will be sorted from large to small according to scores\n keep = batched_nms(boxes, scores, labels, self.nms_thresh)\n\n # Get the top k prediction targets with scores in the top\n keep = keep[:self.detection_per_img]\n boxes, scores, labels = boxes[keep], scores[keep], labels[keep]\n\n all_boxes.append(boxes)\n all_scores.append(scores)\n all_labels.append(labels)\n return all_boxes, all_scores, all_labels\n\n def forward(self, features, proposals, image_shapes, targets=None):\n # type: (Dict[str, Tensor], List[Tensor], List[Tuple[int, int]], Optional[List[Dict[str, Tensor]]])\n \"\"\"\n Arguments:\n features (List[Tensor])\n proposals (List[Tensor[N, 4]])\n image_shapes (List[Tuple[H, W]])\n targets (List[Dict])\n \"\"\"\n\n # Check the data type of targets\n if targets is not None:\n for t in targets:\n floating_point_types = (torch.float, torch.double, torch.half)\n assert t[\"boxes\"].dtype in floating_point_types, \"target boxes must of float type\"\n assert t[\"labels\"].dtype == torch.int64, \"target labels must of int64 type\"\n\n if self.training:\n proposals, matched_idxs, labels, regression_targets = self.select_training_samples(proposals, targets)\n else:\n labels = None\n regression_targets = None\n matched_idxs = None\n\n box_features = self.box_roi_pool(features, proposals, image_shapes)\n box_features = self.box_head(box_features)\n class_logits, box_regression = self.box_predictor(box_features)\n\n result = torch.jit.annotate(List[Dict[str, torch.Tensor]], [])\n losses = {}\n if self.training:\n assert labels is not None and regression_targets is not None\n loss_classifier, loss_box_reg = fastrcnn_loss(\n class_logits, box_regression, labels, regression_targets)\n losses = {\n \"loss_classifier\": loss_classifier,\n \"loss_box_reg\": loss_box_reg\n }\n else:\n boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes)\n num_images = len(boxes)\n for i in range(num_images):\n result.append(\n {\n \"boxes\": boxes[i],\n \"labels\": labels[i],\n \"scores\": scores[i],\n }\n )\n return result, losses","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## RPN","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom torch.jit.annotations import List, Optional, Dict, Tuple\nfrom torch import Tensor\n\n\n@torch.jit.unused\ndef _onnx_get_num_anchors_and_pre_nms_top_n(ob, orig_pre_nms_top_n):\n # type: (Tensor, int) -> Tuple[int, int]\n from torch.onnx import operators\n num_anchors = operators.shape_as_tensor(ob)[1].unsqueeze(0)\n\n pre_nms_top_n = torch.min(torch.cat(\n (torch.tensor([orig_pre_nms_top_n], dtype=num_anchors.dtype),\n num_anchors), 0).to(torch.int32)).to(num_anchors.dtype)\n\n return num_anchors, pre_nms_top_n\n\n\nclass AnchorsGenerator(nn.Module):\n __annotations__ = {\n 'cell_anchors': Optional[List[torch.Tensor]],\n '_cache': Dict[str, List[torch.Tensor]]\n }\n\n def __init__(self, sizes=(128, 256, 512), aspect_ratios=(0.5, 1.0, 2.0)):\n super(AnchorsGenerator, self).__init__()\n\n if not isinstance(sizes[0], (list, tuple)):\n sizes = tuple((s, ) for s in sizes)\n if not isinstance(aspect_ratios[0], (list, tuple)):\n aspect_ratios = (aspect_ratios, ) * len(sizes)\n\n assert len(sizes) == len(aspect_ratios)\n\n self.sizes = sizes\n self.aspect_ratios = aspect_ratios\n self.cell_anchors = None\n self._cache = {}\n\n def generate_anchors(self, scales, aspect_ratios, dtype=torch.float32, device='cpu'):\n # type: (List[int], List[float], int, Device)\n scales = torch.as_tensor(scales, dtype=dtype, device=device)\n aspect_ratios = torch.as_tensor(aspect_ratios, dtype=dtype, device=device)\n h_ratios = torch.sqrt(aspect_ratios)\n w_ratios = 1.0 / h_ratios\n\n ws = (w_ratios[:, None] * scales[None, :]).view(-1)\n hs = (h_ratios[:, None] * scales[None, :]).view(-1)\n\n base_anchors = torch.stack([-ws, -hs, ws, hs], dim=1) / 2\n\n return base_anchors.round()\n\n def set_cell_anchors(self, dtype, device):\n # type: (int, Device) -> None\n if self.cell_anchors is not None:\n cell_anchors = self.cell_anchors\n assert cell_anchors is not None\n\n if cell_anchors[0].device == device:\n return\n\n cell_anchors = [\n self.generate_anchors(sizes, aspect_ratios, dtype, device)\n for sizes, aspect_ratios in zip(self.sizes, self.aspect_ratios)\n ]\n self.cell_anchors = cell_anchors\n\n def num_anchors_per_location(self):\n return [len(s) * len(a) for s, a in zip(self.sizes, self.aspect_ratios)]\n\n def grid_anchors(self, grid_sizes, strides):\n # type: (List[List[int]], List[List[Tensor]])\n anchors = []\n cell_anchors = self.cell_anchors\n assert cell_anchors is not None\n\n for size, stride, base_anchors in zip(grid_sizes, strides, cell_anchors):\n grid_height, grid_width = size\n stride_height, stride_width = stride\n device = base_anchors.device\n\n shifts_x = torch.arange(0, grid_width, dtype=torch.float32, device=device) * stride_width\n shifts_y = torch.arange(0, grid_height, dtype=torch.float32, device=device) * stride_height\n\n shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)\n shift_x = shift_x.reshape(-1)\n shift_y = shift_y.reshape(-1)\n\n shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)\n\n shifts_anchor = shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)\n anchors.append(shifts_anchor.reshape(-1, 4))\n\n return anchors\n\n def cached_grid_anchors(self, grid_sizes, strides):\n # type: (List[List[int]], List[List[Tensor]])\n \"\"\"Cache all the calculated anchor information\"\"\"\n key = str(grid_sizes) + str(strides)\n\n if key in self._cache:\n return self._cache[key]\n anchors = self.grid_anchors(grid_sizes, strides)\n self._cache[key] = anchors\n return anchors\n\n def forward(self, image_list, feature_maps):\n # type: (ImageList, List[Tensor])\n\n grid_sizes = list([feature_map.shape[-2:] for feature_map in feature_maps])\n\n image_size = image_list.tensors.shape[-2:]\n dtype, device = feature_maps[0].dtype, feature_maps[0].device\n\n strides = [[torch.tensor(image_size[0] / g[0], dtype=torch.int64, device=device),\n torch.tensor(image_size[1] / g[1], dtype=torch.int64, device=device)] for g in grid_sizes]\n\n # Generate anchors template\n self.set_cell_anchors(dtype, device)\n\n anchors_over_all_feature_maps = self.cached_grid_anchors(grid_sizes, strides)\n\n anchors = torch.jit.annotate(List[List[torch.Tensor]], [])\n\n for i, (image_height, image_width) in enumerate(image_list.image_sizes):\n anchors_in_image = []\n for anchors_per_feature_map in anchors_over_all_feature_maps:\n anchors_in_image.append(anchors_per_feature_map)\n anchors.append(anchors_in_image)\n anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors]\n self._cache.clear()\n return anchors\n\n\nclass RPNHead(nn.Module):\n def __init__(self, in_channels, num_anchors):\n super(RPNHead, self).__init__()\n self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)\n # Calculate the predicted target probability\n self.cls_logits = nn.Conv2d(in_channels, num_anchors, kernel_size=1, stride=1)\n self.bbox_pred = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=1, stride=1)\n\n for layer in self.children():\n if isinstance(layer, nn.Conv2d):\n nn.init.normal_(layer.weight, std=0.01)\n nn.init.constant_(layer.bias, 0)\n\n def forward(self, x):\n # type: (List[Tensor])\n logits = []\n bbox_reg = []\n for i, feature in enumerate(x):\n t = F.relu(self.conv(feature))\n logits.append(self.cls_logits(t))\n bbox_reg.append(self.bbox_pred(t))\n return logits, bbox_reg\n\n\ndef permute_and_flatten(layer, N, A, C, H, W):\n # type: (Tensor, int, int, int, int, int)\n \"\"\"Adjust the tensor order and reshape\"\"\"\n layer = layer.view(N, -1, C, H, W)\n layer = layer.permute(0, 3, 4, 1, 2) # [N, H, W, -1, C]\n layer = layer.reshape(N, -1, C)\n return layer\n\n\ndef concat_box_prediction_layers(box_cls, box_regression):\n # type: (List[Tensor], List[Tensor])\n box_cls_flattened = []\n box_regression_flattened = []\n\n for box_cls_per_level, box_regression_per_level in zip(box_cls, box_regression):\n # [batch_size, anchors_num_per_position * classes_num, height, width]\n N, AxC, H, W = box_cls_per_level.shape\n Ax4 = box_regression_per_level.shape[1]\n A = Ax4 // 4\n # classes_num\n C = AxC // A\n\n box_cls_per_level = permute_and_flatten(box_cls_per_level, N, A, C, H, W)\n box_cls_flattened.append(box_cls_per_level)\n\n box_regression_per_level = permute_and_flatten(box_regression_per_level, N, A, 4, H, W)\n box_regression_flattened.append(box_regression_per_level)\n box_cls = torch.cat(box_cls_flattened, dim=1).flatten(0, -2)\n box_regression = torch.cat(box_regression_flattened, dim=1).reshape(-1, 4)\n return box_cls, box_regression\n\n\nclass RegionProposalNetwork(nn.Module):\n __annotations__ = {\n 'box_coder': BoxCoder,\n 'proposal_matcher': Matcher,\n 'fg_bg_sampler': BalancedPositiveNegativeSampler,\n 'pre_nms_top_n': Dict[str, int],\n 'post_nms_top_n': Dict[str, int],\n }\n\n def __init__(self, anchor_generator, head, fg_iou_thresh, bg_iou_thresh,\n batch_size_per_image, positive_fraction, pre_nms_top_n, post_nms_top_n, nms_thresh):\n super(RegionProposalNetwork, self).__init__()\n self.anchor_generator = anchor_generator\n self.head = head\n self.box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0))\n\n self.box_similarity = box_iou\n\n self.proposal_matcher = Matcher(\n fg_iou_thresh,\n bg_iou_thresh,\n allow_low_quality_matches=True\n )\n\n self.fg_bg_sampler = BalancedPositiveNegativeSampler(\n batch_size_per_image, positive_fraction # 256, 0.5\n )\n\n # use during testing\n self._pre_nms_top_n = pre_nms_top_n\n self._post_nms_top_n = post_nms_top_n\n self.nms_thresh = nms_thresh\n self.min_size = 1e-3\n\n def pre_nms_top_n(self):\n if self.training:\n return self._pre_nms_top_n['training']\n return self._pre_nms_top_n['testing']\n\n def post_nms_top_n(self):\n if self.training:\n return self._post_nms_top_n['training']\n return self._post_nms_top_n['testing']\n\n def assign_targets_to_anchors(self, anchors, targets):\n # type: (List[Tensor], List[Dict[str, Tensor]])\n labels = []\n matched_gt_boxes = []\n for anchors_per_image, targets_per_image in zip(anchors, targets):\n gt_boxes = targets_per_image[\"boxes\"]\n if gt_boxes.numel() == 0:\n device = anchors_per_image.device\n matched_gt_boxes_per_image = torch.zeros(anchors_per_image.shape, dtype=torch.float32, device=device)\n labels_per_image = torch.zeros((anchors_per_image.shape[0],), dtype=torch.float32, device=device)\n else:\n # Calculate the IOU information of anchors and real bbox\n match_quality_matrix = box_iou(gt_boxes, anchors_per_image)\n matched_idxs = self.proposal_matcher(match_quality_matrix)\n\n matched_gt_boxes_per_image = gt_boxes[matched_idxs.clamp(min=0)]\n\n labels_per_image = matched_idxs >= 0\n labels_per_image = labels_per_image.to(dtype=torch.float32)\n\n # background (negative examples)\n bg_indices = matched_idxs == self.proposal_matcher.BELOW_LOW_THRESHOLD\n labels_per_image[bg_indices] = torch.tensor(0.0)\n\n # discard indices that are between thresholds\n inds_to_discard = matched_idxs == self.proposal_matcher.BETWEEN_THRESHOLDS # -2\n labels_per_image[inds_to_discard] = torch.tensor(-1.0)\n\n labels.append(labels_per_image)\n matched_gt_boxes.append(matched_gt_boxes_per_image)\n return labels, matched_gt_boxes\n\n def _get_top_n_idx(self, objectness, num_anchors_per_level):\n # type: (Tensor, List[int])\n r = [] \n offset = 0\n for ob in objectness.split(num_anchors_per_level, 1):\n if torchvision._is_tracing():\n num_anchors, pre_nms_top_n = _onnx_get_num_anchors_and_pre_nms_top_n(ob, self.pre_nms_top_n())\n else:\n num_anchors = ob.shape[1]\n pre_nms_top_n = min(self.pre_nms_top_n(), num_anchors) # self.pre_nms_top_n=1000\n # Returns the k largest elements of the given input tensor along a given dimension\n _, top_n_idx = ob.topk(pre_nms_top_n, dim=1)\n r.append(top_n_idx + offset)\n offset += num_anchors\n return torch.cat(r, dim=1)\n\n def filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_level):\n # type: (Tensor, Tensor, List[Tuple[int, int]], List[int])\n num_images = proposals.shape[0]\n device = proposals.device\n\n objectness = objectness.detach()\n objectness = objectness.reshape(num_images, -1)\n\n levels = [torch.full((n, ), idx, dtype=torch.int64, device=device)\n for idx, n in enumerate(num_anchors_per_level)]\n levels = torch.cat(levels, 0)\n\n # Expand this tensor to the same size as objectness\n levels = levels.reshape(1, -1).expand_as(objectness)\n\n top_n_idx = self._get_top_n_idx(objectness, num_anchors_per_level)\n\n image_range = torch.arange(num_images, device=device)\n batch_idx = image_range[:, None] # [batch_size, 1]\n\n objectness = objectness[batch_idx, top_n_idx]\n levels = levels[batch_idx, top_n_idx]\n proposals = proposals[batch_idx, top_n_idx]\n\n final_boxes = []\n final_scores = []\n for boxes, scores, lvl, img_shape in zip(proposals, objectness, levels, image_shapes):\n boxes = clip_boxes_to_image(boxes, img_shape)\n keep = remove_small_boxes(boxes, self.min_size)\n boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep]\n # non-maximum suppression, independently done per level\n keep = batched_nms(boxes, scores, lvl, self.nms_thresh)\n # keep only topk scoring predictions\n keep = keep[: self.post_nms_top_n()]\n boxes, scores = boxes[keep], scores[keep]\n final_boxes.append(boxes)\n final_scores.append(scores)\n return final_boxes, final_scores\n\n def compute_loss(self, objectness, pred_bbox_deltas, labels, regression_targets):\n # type: (Tensor, Tensor, List[Tensor], List[Tensor])\n \"\"\"\n Calculate RPN loss\n \"\"\"\n sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)\n sampled_pos_inds = torch.nonzero(torch.cat(sampled_pos_inds, dim=0)).squeeze(1)\n sampled_neg_inds = torch.nonzero(torch.cat(sampled_neg_inds, dim=0)).squeeze(1)\n\n sampled_inds = torch.cat([sampled_pos_inds, sampled_neg_inds], dim=0)\n objectness = objectness.flatten()\n\n labels = torch.cat(labels, dim=0)\n regression_targets = torch.cat(regression_targets, dim=0)\n\n box_loss = smooth_l1_loss(\n pred_bbox_deltas[sampled_pos_inds],\n regression_targets[sampled_pos_inds],\n beta=1 / 9,\n size_average=False,\n ) / (sampled_inds.numel())\n\n objectness_loss = F.binary_cross_entropy_with_logits(\n objectness[sampled_inds], labels[sampled_inds]\n )\n\n return objectness_loss, box_loss\n\n def forward(self, images, features, targets=None):\n # type: (ImageList, Dict[str, Tensor], Optional[List[Dict[str, Tensor]]])\n features = list(features.values())\n\n objectness, pred_bbox_deltas = self.head(features)\n\n anchors = self.anchor_generator(images, features)\n\n # batch_size\n num_images = len(anchors)\n\n # numel() Returns the total number of elements in the input tensor.\n num_anchors_per_level_shape_tensors = [o[0].shape for o in objectness]\n num_anchors_per_level = [s[0] * s[1] * s[2] for s in num_anchors_per_level_shape_tensors]\n\n objectness, pred_bbox_deltas = concat_box_prediction_layers(objectness, pred_bbox_deltas)\n\n proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors)\n proposals = proposals.view(num_images, -1, 4)\n\n boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)\n\n losses = {}\n if self.training:\n assert targets is not None\n labels, matched_gt_boxes = self.assign_targets_to_anchors(anchors, targets)\n regression_targets = self.box_coder.encode(matched_gt_boxes, anchors)\n loss_objectness, loss_rpn_box_reg = self.compute_loss(\n objectness, pred_bbox_deltas, labels, regression_targets\n )\n losses = {\n \"loss_objectness\": loss_objectness,\n \"loss_rpn_box_reg\": loss_rpn_box_reg\n }\n return boxes, losses\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Faster R-CNN","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch import Tensor\nimport torch.nn.functional as F\nfrom torchvision.ops import MultiScaleRoIAlign\nfrom torch.jit.annotations import Tuple, List, Dict, Optional\nfrom collections import OrderedDict\nimport warnings\n\n\nclass FasterRCNNBase(nn.Module):\n def __init__(self, backbone, rpn, roi_heads, transform):\n super(FasterRCNNBase, self).__init__()\n self.backbone = backbone\n self.rpn = rpn\n self.roi_heads = roi_heads\n self.transform = transform\n self._has_warned = False\n\n @torch.jit.unused\n def eager_outputs(self, losses, detections):\n # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]\n if self.training:\n return losses\n\n return detections\n\n def forward(self, images, targets=None):\n # type: (List[Tensor], Optional[List[Dict[str, Tensor]]])\n if self.training and targets is None:\n raise ValueError(\"In training mode, targets should be passed\")\n\n if self.training:\n assert targets is not None\n for target in targets:\n boxes = target[\"boxes\"]\n if isinstance(boxes, torch.Tensor):\n if len(boxes.shape) != 2 or boxes.shape[-1] != 4:\n raise ValueError(\"Expected target boxes to be a tensor of shape [N, 4], got {:}.\".format(boxes.shape))\n else:\n raise ValueError(\"Expected target boxes to be of type Tensor, got {:}.\".format(type(boxes)))\n\n original_image_sizes = torch.jit.annotate(List[Tuple[int, int]], [])\n for img in images:\n val = img.shape[-2:]\n assert len(val) == 2\n original_image_sizes.append((val[0], val[1]))\n\n images, targets = self.transform(images, targets)\n features = self.backbone(images.tensors)\n if isinstance(features, torch.Tensor):\n features = OrderedDict([('0', features)])\n\n proposals, proposal_losses = self.rpn(images, features, targets)\n\n detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)\n\n detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)\n\n losses = {}\n losses.update(detector_losses)\n losses.update(proposal_losses)\n\n if torch.jit.is_scripting():\n if not self._has_warned:\n warnings.warn(\"RCNN always returns a (Losses, Detections) tuple in scripting\")\n self._has_warned = True\n return losses, detections\n else:\n return self.eager_outputs(losses, detections)\n\n\nclass TwoMLPHead(nn.Module):\n def __init__(self, in_channels, representation_size):\n super(TwoMLPHead, self).__init__()\n\n self.fc6 = nn.Linear(in_channels, representation_size)\n self.fc7 = nn.Linear(representation_size, representation_size)\n\n def forward(self, x):\n x = x.flatten(start_dim=1)\n\n x = F.relu(self.fc6(x))\n x = F.relu(self.fc7(x))\n\n return x\n\n\nclass FastRCNNPredictor(nn.Module):\n def __init__(self, in_channels, num_classes):\n super(FastRCNNPredictor, self).__init__()\n\n self.cls_score = nn.Linear(in_channels, num_classes)\n self.bbox_pred = nn.Linear(in_channels, num_classes * 4)\n\n def forward(self, x):\n if x.dim() == 4:\n assert list(x.shape[2:]) == [1, 1]\n x = x.flatten(start_dim=1)\n scores = self.cls_score(x)\n bbox_deltas = self.bbox_pred(x)\n\n return scores, bbox_deltas\n\n\nclass FasterRCNN(FasterRCNNBase):\n def __init__(self, backbone, num_classes=None,\n # transform parameter\n min_size=800, max_size=1344,\n image_mean=None, image_std=None,\n # RPN parameters\n rpn_anchor_generator=None, rpn_head=None,\n rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000,\n rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000,\n rpn_nms_thresh=0.7,\n rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3,\n rpn_batch_size_per_image=256, rpn_positive_fraction=0.5,\n # Box parameters\n box_roi_pool=None, box_head=None, box_predictor=None,\n box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,\n box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5,\n box_batch_size_per_image=512, box_positive_fraction=0.25,\n bbox_reg_weights=None):\n if not hasattr(backbone, \"out_channels\"):\n raise ValueError(\n \"backbone should contain an attribute out_channels\"\n \"specifying the number of output channels (assumed to be the\"\n \"same for all the levels\"\n )\n\n assert isinstance(rpn_anchor_generator, (AnchorsGenerator, type(None)))\n assert isinstance(box_roi_pool, (MultiScaleRoIAlign, type(None)))\n\n if num_classes is not None:\n if box_predictor is not None:\n raise ValueError(\"num_classes should be None when box_predictor is specified\")\n else:\n if box_predictor is None:\n raise ValueError(\"num_classes should not be None when box_predictor is not specified\")\n\n out_channels = backbone.out_channels\n\n if rpn_anchor_generator is None:\n anchor_sizes = ((32,), (64,), (128,), (256,), (512,))\n aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)\n rpn_anchor_generator = AnchorsGenerator(\n anchor_sizes, aspect_ratios\n )\n\n if rpn_head is None:\n rpn_head = RPNHead(\n out_channels, rpn_anchor_generator.num_anchors_per_location()[0]\n )\n\n rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)\n rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)\n\n rpn = RegionProposalNetwork(\n rpn_anchor_generator, rpn_head,\n rpn_fg_iou_thresh, rpn_bg_iou_thresh,\n rpn_batch_size_per_image, rpn_positive_fraction,\n rpn_pre_nms_top_n, rpn_post_nms_top_n, rpn_nms_thresh)\n\n if box_roi_pool is None:\n box_roi_pool = MultiScaleRoIAlign(\n featmap_names=['0', '1', '2', '3'],\n output_size=[7, 7],\n sampling_ratio=2)\n\n if box_head is None:\n resolution = box_roi_pool.output_size[0]\n representation_size = 1024\n box_head = TwoMLPHead(\n out_channels * resolution ** 2,\n representation_size\n )\n\n if box_predictor is None:\n representation_size = 1024\n box_predictor = FastRCNNPredictor(\n representation_size,\n num_classes)\n\n roi_heads = ROIHeads(\n # box\n box_roi_pool, box_head, box_predictor,\n box_fg_iou_thresh, box_bg_iou_thresh,\n box_batch_size_per_image, box_positive_fraction,\n bbox_reg_weights,\n box_score_thresh, box_nms_thresh, box_detections_per_img)\n\n if image_mean is None:\n image_mean = [0.485, 0.456, 0.406]\n if image_std is None:\n image_std = [0.229, 0.224, 0.225]\n\n transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std)\n\n super(FasterRCNN, self).__init__(backbone, rpn, roi_heads, transform)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_model(num_classes):\n backbone = resnet50_fpn_backbone()\n model = FasterRCNN(backbone=backbone, num_classes=91)\n \n weights_dict = torch.load(\"../input/weights/5-resNetFpn-model-19.pth\")\n missing_keys, unexpected_keys = model.load_state_dict(weights_dict, strict=False)\n if len(missing_keys) != 0 or len(unexpected_keys) != 0:\n print(\"missing_keys: \", missing_keys)\n print(\"unexpected_keys: \", unexpected_keys)\n\n # get 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 return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport os\nimport time\nimport warnings\n\nwarnings.filterwarnings('ignore')\n\n\ndef collate_fn(batch):\n return tuple(zip(*batch))\n\n\ndef main():\n since = time.time()\n device = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\n print(device)\n\n data_transform = {\n \"train\": transforms.Compose([transforms.ToTensor(),\n transforms.RandomHorizontalFlip(0.5)]),\n \"val\": transforms.Compose([transforms.ToTensor()])\n }\n\n data_root = \"../input/global-wheat-detection\"\n # load train data set\n train_data_set = WheatDataSet(data_root, data_transform[\"train\"])\n train_data_loader = torch.utils.data.DataLoader(train_data_set,\n batch_size=4,\n shuffle=True,\n num_workers=2,\n collate_fn=collate_fn)\n\n # create model num_classes equal background + 20 classes\n model = create_model(num_classes=2)\n\n model.to(device)\n\n # define optimizer\n params = [p for p in model.parameters() if p.requires_grad]\n optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n # learning rate scheduler\n lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,\n step_size=5,\n gamma=0.33)\n # You can install pycocotools to simplify your training.I omitted these steps","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prediction","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nfrom torchvision import transforms\nfrom PIL import Image\nimport json\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport time\nimport warnings\nimport glob\n\nwarnings.filterwarnings('ignore')\n\n\ndef create_model(num_classes):\n backbone = resnet50_fpn_backbone()\n model = FasterRCNN(backbone=backbone, num_classes=num_classes)\n\n return model\n\n\n# get devices\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)\n\n# create model\nmodel = create_model(num_classes=2)\n\n# load train weights\ntrain_weights = \"../input/weights/15-fasterrcnn_resnet50_fpn_best.pth\"\n# model.load_state_dict(torch.load(train_weights)[\"model\"])\nmodel.load_state_dict(torch.load(train_weights))\nmodel.to(device)\n\n# read class_indict\ncategory_index = {\"wheat\": 1}\n\n\norigin_list = glob.glob('../input/global-wheat-detection/test/*.jpg')\nprint(origin_list)\n\npreds = []\n\nfor i, img_name in enumerate(origin_list):\n # load image\n original_img = Image.open(img_name)\n\n # from pil image to tensor, do not normalize image\n data_transform = transforms.Compose([transforms.ToTensor()])\n img = data_transform(original_img)\n # expand batch dimension\n img = torch.unsqueeze(img, dim=0)\n\n model.eval()\n with torch.no_grad():\n since = time.time()\n predictions = model(img.to(device))[0]\n print('{} Time:{}s'.format(i, time.time() - since))\n predict_boxes = predictions[\"boxes\"].to(\"cpu\").numpy()\n predict_classes = predictions[\"labels\"].to(\"cpu\").numpy()\n predict_scores = predictions[\"scores\"].to(\"cpu\").numpy()\n\n draw_box(original_img,\n predict_boxes,\n predict_classes,\n predict_scores,\n category_index,\n thresh=0.5,\n line_thickness=5)\n plt.imshow(original_img)\n plt.show()\n\n predict = \"\"\n for box, score in zip(predict_boxes, predict_scores):\n str_box = \"\"\n box[2] = box[2] - box[0]\n box[3] = box[3] - box[1]\n for b in box:\n str_box += str(b) + ' '\n predict += str(score) + ' ' + str_box\n preds.append(predict)\n\nname_list = [name.split('/')[-1].split('.')[0] for name in origin_list]\nprint(name_list)\ndataframe = pd.DataFrame({\"image_id\": name_list, \"PredictionString\": preds})\ndataframe.to_csv('submission.csv', index=False)","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}