{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import numpy as np\nimport cv2\nimport random\n\n\ndef flip(img):\n return img[:, :, ::-1].copy()\n\n# todo what the hell is this?\ndef get_border(border, size):\n i = 1\n while size - border // i <= border // i:\n i *= 2\n return border // i\n\ndef transform_preds(coords, center, scale, output_size):\n target_coords = np.zeros(coords.shape)\n trans = get_affine_transform(center, scale, 0, output_size, inv=1)\n for p in range(coords.shape[0]):\n target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)\n return target_coords\n\n\ndef get_affine_transform(center,\n scale,\n rot,\n output_size,\n shift=np.array([0, 0], dtype=np.float32),\n inv=0):\n if not isinstance(scale, np.ndarray) and not isinstance(scale, list):\n scale = np.array([scale, scale], dtype=np.float32)\n\n scale_tmp = scale\n src_w = scale_tmp[0]\n dst_w = output_size[0]\n dst_h = output_size[1]\n\n rot_rad = np.pi * rot / 180\n src_dir = get_dir([0, src_w * -0.5], rot_rad)\n dst_dir = np.array([0, dst_w * -0.5], np.float32)\n\n src = np.zeros((3, 2), dtype=np.float32)\n dst = np.zeros((3, 2), dtype=np.float32)\n src[0, :] = center + scale_tmp * shift\n src[1, :] = center + src_dir + scale_tmp * shift\n dst[0, :] = [dst_w * 0.5, dst_h * 0.5]\n dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir\n\n src[2:, :] = get_3rd_point(src[0, :], src[1, :])\n dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])\n\n if inv:\n trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))\n else:\n trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))\n\n return trans\n\n\ndef affine_transform(pt, t):\n new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T\n new_pt = np.dot(t, new_pt)\n return new_pt[:2]\n\n\ndef get_3rd_point(a, b):\n direct = a - b\n return b + np.array([-direct[1], direct[0]], dtype=np.float32)\n\n\ndef get_dir(src_point, rot_rad):\n _sin, _cos = np.sin(rot_rad), np.cos(rot_rad)\n\n src_result = [0, 0]\n src_result[0] = src_point[0] * _cos - src_point[1] * _sin\n src_result[1] = src_point[0] * _sin + src_point[1] * _cos\n\n return src_result\n\n\ndef crop(img, center, scale, output_size, rot=0):\n trans = get_affine_transform(center, scale, rot, output_size)\n\n dst_img = cv2.warpAffine(img,\n trans,\n (int(output_size[0]), int(output_size[1])),\n flags=cv2.INTER_LINEAR)\n\n return dst_img\n\n\ndef gaussian_radius(det_size, min_overlap=0.7):\n height, width = det_size\n\n a1 = 1\n b1 = (height + width)\n c1 = width * height * (1 - min_overlap) / (1 + min_overlap)\n sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)\n # r1 = (b1 + sq1) / 2 #\n r1 = (b1 - sq1) / (2 * a1)\n\n a2 = 4\n b2 = 2 * (height + width)\n c2 = (1 - min_overlap) * width * height\n sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)\n # r2 = (b2 + sq2) / 2\n r2 = (b2 - sq2) / (2 * a2)\n\n a3 = 4 * min_overlap\n b3 = -2 * min_overlap * (height + width)\n c3 = (min_overlap - 1) * width * height\n sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)\n r3 = (b3 + sq3) / 2\n # r3 = (b3 + sq3) / (2 * a3)\n return min(r1, r2, r3)\n\n\ndef gaussian2D(shape, sigma=1):\n m, n = [(ss - 1.) / 2. for ss in shape]\n y, x = np.ogrid[-m:m + 1, -n:n + 1]\n\n h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))\n h[h < np.finfo(h.dtype).eps * h.max()] = 0\n return h\n\n\ndef draw_umich_gaussian(heatmap, center, radius, k=1):\n diameter = 2 * radius + 1\n gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)\n\n x, y = int(center[0]), int(center[1])\n\n height, width = heatmap.shape[0:2]\n\n left, right = min(x, radius), min(width - x, radius + 1)\n top, bottom = min(y, radius), min(height - y, radius + 1)\n\n masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]\n masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]\n if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug\n np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)\n return heatmap\n\n\ndef draw_dense_reg(regmap, heatmap, center, value, radius, is_offset=False):\n diameter = 2 * radius + 1\n gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)\n value = np.array(value, dtype=np.float32).reshape(-1, 1, 1)\n dim = value.shape[0]\n reg = np.ones((dim, diameter * 2 + 1, diameter * 2 + 1), dtype=np.float32) * value\n if is_offset and dim == 2:\n delta = np.arange(diameter * 2 + 1) - radius\n reg[0] = reg[0] - delta.reshape(1, -1)\n reg[1] = reg[1] - delta.reshape(-1, 1)\n\n x, y = int(center[0]), int(center[1])\n\n height, width = heatmap.shape[0:2]\n\n left, right = min(x, radius), min(width - x, radius + 1)\n top, bottom = min(y, radius), min(height - y, radius + 1)\n\n masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]\n masked_regmap = regmap[:, y - top:y + bottom, x - left:x + right]\n masked_gaussian = gaussian[radius - top:radius + bottom,\n radius - left:radius + right]\n masked_reg = reg[:, radius - top:radius + bottom,\n radius - left:radius + right]\n if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug\n idx = (masked_gaussian >= masked_heatmap).reshape(\n 1, masked_gaussian.shape[0], masked_gaussian.shape[1])\n masked_regmap = (1 - idx) * masked_regmap + idx * masked_reg\n regmap[:, y - top:y + bottom, x - left:x + right] = masked_regmap\n return regmap\n\n\ndef draw_msra_gaussian(heatmap, center, sigma):\n tmp_size = sigma * 3\n mu_x = int(center[0] + 0.5)\n mu_y = int(center[1] + 0.5)\n w, h = heatmap.shape[0], heatmap.shape[1]\n ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]\n br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]\n if ul[0] >= h or ul[1] >= w or br[0] < 0 or br[1] < 0:\n return heatmap\n size = 2 * tmp_size + 1\n x = np.arange(0, size, 1, np.float32)\n y = x[:, np.newaxis]\n x0 = y0 = size // 2\n g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))\n g_x = max(0, -ul[0]), min(br[0], h) - ul[0]\n g_y = max(0, -ul[1]), min(br[1], w) - ul[1]\n img_x = max(0, ul[0]), min(br[0], h)\n img_y = max(0, ul[1]), min(br[1], w)\n heatmap[img_y[0]:img_y[1], img_x[0]:img_x[1]] = np.maximum(\n heatmap[img_y[0]:img_y[1], img_x[0]:img_x[1]],\n g[g_y[0]:g_y[1], g_x[0]:g_x[1]])\n return heatmap\n\n\ndef grayscale(image):\n return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n\n\ndef lighting_(data_rng, image, alphastd, eigval, eigvec):\n alpha = data_rng.normal(scale=alphastd, size=(3,))\n image += np.dot(eigvec, eigval * alpha)\n\n\ndef blend_(alpha, image1, image2):\n image1 *= alpha\n image2 *= (1 - alpha)\n image1 += image2\n\n\ndef saturation_(data_rng, image, gs, gs_mean, var):\n alpha = 1. + data_rng.uniform(low=-var, high=var)\n blend_(alpha, image, gs[:, :, None])\n\n\ndef brightness_(data_rng, image, gs, gs_mean, var):\n alpha = 1. + data_rng.uniform(low=-var, high=var)\n image *= alpha\n\n\ndef contrast_(data_rng, image, gs, gs_mean, var):\n alpha = 1. + data_rng.uniform(low=-var, high=var)\n blend_(alpha, image, gs_mean)\n\n\ndef color_aug(data_rng, image, eig_val, eig_vec):\n functions = [brightness_, contrast_, saturation_]\n random.shuffle(functions)\n\n gs = grayscale(image)\n gs_mean = gs.mean()\n for f in functions:\n f(data_rng, image, gs, gs_mean, 0.4)\n lighting_(data_rng, image, 0.1, eig_val, eig_vec)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define model -- using resnet18 as backbone","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# import torch\n# import torch.nn as nn\n# import torch.utils.model_zoo as model_zoo\n\n# BN_MOMENTUM = 0.1\n\n# model_urls = {'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n# 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n# 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n# 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n# 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', }\n\n\n# def conv3x3(in_planes, out_planes, stride=1):\n# \"\"\"3x3 convolution with padding\"\"\"\n# return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n\n\n# class BasicBlock(nn.Module):\n# expansion = 1\n\n# def __init__(self, inplanes, planes, stride=1, downsample=None):\n# super(BasicBlock, self).__init__()\n# self.conv1 = conv3x3(inplanes, planes, stride)\n# self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.relu = nn.ReLU(inplace=True)\n# self.conv2 = conv3x3(planes, planes)\n# self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.downsample = downsample\n# self.stride = stride\n\n# def forward(self, x):\n# residual = 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\n# if self.downsample is not None:\n# residual = self.downsample(x)\n\n# out += residual\n# out = self.relu(out)\n\n# return out\n\n\n# class Bottleneck(nn.Module):\n# expansion = 4\n\n# def __init__(self, inplanes, planes, stride=1, downsample=None):\n# super(Bottleneck, self).__init__()\n# self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n# self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n# self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n# self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM)\n# self.relu = nn.ReLU(inplace=True)\n# self.downsample = downsample\n# self.stride = stride\n\n# def forward(self, x):\n# residual = 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# if self.downsample is not None:\n# residual = self.downsample(x)\n\n# out += residual\n# out = self.relu(out)\n\n# return out\n\n\n# class PoseResNet(nn.Module):\n# def __init__(self, block, layers, head_conv, num_classes):\n# super(PoseResNet, self).__init__()\n# self.inplanes = 64\n# self.deconv_with_bias = False\n# self.num_classes = num_classes\n\n# self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n# self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\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, layers[0])\n# self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n# self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n# self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n\n# # used for deconv layers\n# self.deconv_layers = self._make_deconv_layer(3, [256, 256, 256], [4, 4, 4])\n# # self.final_layer = []\n\n# if head_conv > 0:\n# # heatmap layers\n# self.hmap = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, num_classes, kernel_size=1))\n# # self.hmap[-1].bias.data.fill_(-2.19)\n# # regression layers\n# self.regs = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, 2, kernel_size=1))\n# self.w_h_ = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, 2, kernel_size=1))\n# else:\n# # heatmap layers\n# self.hmap = nn.Conv2d(in_channels=256, out_channels=num_classes, kernel_size=1)\n# # regression layers\n# self.regs = nn.Conv2d(in_channels=256, out_channels=2, kernel_size=1)\n# self.w_h_ = nn.Conv2d(in_channels=256, out_channels=2, kernel_size=1)\n\n# # self.final_layer = nn.ModuleList(self.final_layer)\n\n# def _make_layer(self, block, planes, blocks, stride=1):\n# downsample = None\n# if stride != 1 or self.inplanes != planes * block.expansion:\n# downsample = nn.Sequential(nn.Conv2d(self.inplanes, planes * block.expansion,\n# kernel_size=1, stride=stride, bias=False),\n# nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM))\n\n# layers = []\n# layers.append(block(self.inplanes, planes, stride, downsample))\n# self.inplanes = planes * block.expansion\n# for i in range(1, blocks):\n# layers.append(block(self.inplanes, planes))\n# return nn.Sequential(*layers)\n\n# def _get_deconv_cfg(self, deconv_kernel, index):\n# if deconv_kernel == 4:\n# padding = 1\n# output_padding = 0\n# elif deconv_kernel == 3:\n# padding = 1\n# output_padding = 1\n# elif deconv_kernel == 2:\n# padding = 0\n# output_padding = 0\n\n# return deconv_kernel, padding, output_padding\n\n# def _make_deconv_layer(self, num_layers, num_filters, num_kernels):\n# assert num_layers == len(num_filters), \\\n# 'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n# assert num_layers == len(num_kernels), \\\n# 'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n\n# layers = []\n# for i in range(num_layers):\n# kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i], i)\n\n# planes = num_filters[i]\n# layers.append(nn.ConvTranspose2d(in_channels=self.inplanes,\n# out_channels=planes,\n# kernel_size=kernel,\n# stride=2,\n# padding=padding,\n# output_padding=output_padding,\n# bias=self.deconv_with_bias))\n# layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))\n# layers.append(nn.ReLU(inplace=True))\n# self.inplanes = planes\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# x = self.deconv_layers(x)\n# out = [[self.hmap(x), self.regs(x), self.w_h_(x)]]\n# return out\n\n# def init_weights(self, num_layers):\n# for m in self.deconv_layers.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n# for m in self.hmap.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.constant_(m.bias, -2.19)\n# for m in self.regs.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.normal_(m.weight, std=0.001)\n# nn.init.constant_(m.bias, 0)\n# for m in self.w_h_.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.normal_(m.weight, std=0.001)\n# nn.init.constant_(m.bias, 0)\n# url = model_urls['resnet{}'.format(num_layers)]\n# pretrained_state_dict = model_zoo.load_url(url)\n# print('=> loading pretrained model {}'.format(url))\n# self.load_state_dict(pretrained_state_dict, strict=False)\n\n\n# resnet_spec = {18: (BasicBlock, [2, 2, 2, 2]),\n# 34: (BasicBlock, [3, 4, 6, 3]),\n# 50: (Bottleneck, [3, 4, 6, 3]),\n# 101: (Bottleneck, [3, 4, 23, 3]),\n# 152: (Bottleneck, [3, 8, 36, 3])}\n\n# def resnet_18():\n# model = PoseResNet(BasicBlock, [2, 2, 2, 2], head_conv=64, num_classes=80)\n# model.init_weights(18)\n# return model\n\n# def get_pose_net(num_layers, head_conv, num_classes=80):\n# block_class, layers = resnet_spec[num_layers]\n\n# model = PoseResNet(block_class, layers, head_conv=head_conv, num_classes=num_classes)\n# # model.init_weights(num_layers)\n# return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# import torch\n# import torch.nn as nn\n# import torch.utils.model_zoo as model_zoo\n\n# BN_MOMENTUM = 0.1\n\n# model_urls = {'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n# 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n# 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n# 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n# 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', }\n\n\n# def conv3x3(in_planes, out_planes, stride=1):\n# \"\"\"3x3 convolution with padding\"\"\"\n# return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n\n\n# class BasicBlock(nn.Module):\n# expansion = 1\n\n# def __init__(self, inplanes, planes, stride=1, downsample=None):\n# super(BasicBlock, self).__init__()\n# self.conv1 = conv3x3(inplanes, planes, stride)\n# self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.relu = nn.ReLU(inplace=True)\n# self.conv2 = conv3x3(planes, planes)\n# self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.downsample = downsample\n# self.stride = stride\n\n# def forward(self, x):\n# residual = 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\n# if self.downsample is not None:\n# residual = self.downsample(x)\n\n# out += residual\n# out = self.relu(out)\n\n# return out\n\n\n# class Bottleneck(nn.Module):\n# expansion = 4\n\n# def __init__(self, inplanes, planes, stride=1, downsample=None):\n# super(Bottleneck, self).__init__()\n# self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n# self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n# self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n# self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM)\n# self.relu = nn.ReLU(inplace=True)\n# self.downsample = downsample\n# self.stride = stride\n\n# def forward(self, x):\n# residual = 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# if self.downsample is not None:\n# residual = self.downsample(x)\n\n# out += residual\n# out = self.relu(out)\n\n# return out\n\n\n# class PoseResNet(nn.Module):\n# def __init__(self, block, layers, head_conv, num_classes, verbose = True):\n# super(PoseResNet, self).__init__()\n# self.inplanes = 64\n# self.deconv_with_bias = False\n# self.num_classes = num_classes\n# self.verbose = verbose\n\n# self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n# self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\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, layers[0])\n# self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n# self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n# self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n\n# # used for deconv layers\n# self.deconv_layer16 = self._make_deconv_layer(3, [128, 128, 128], [4, 4, 4],2048)\n# self.deconv_layer32 = self._make_deconv_layer(2, [64,64], [4, 4],1024)\n# self.deconv_layer64 = self._make_deconv_layer(1, [64], [4],512)\n# # self.final_layer = []\n\n# if head_conv > 0:\n# # heatmap layers\n# self.hmap = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, num_classes, kernel_size=1))\n# # self.hmap[-1].bias.data.fill_(-2.19)\n# # regression layers\n# self.regs = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, 2, kernel_size=1))\n# self.w_h_ = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, 2, kernel_size=1))\n# else:\n# # heatmap layers\n# self.hmap = nn.Conv2d(in_channels=256, out_channels=num_classes, kernel_size=1)\n# # regression layers\n# self.regs = nn.Conv2d(in_channels=256, out_channels=2, kernel_size=1)\n# self.w_h_ = nn.Conv2d(in_channels=256, out_channels=2, kernel_size=1)\n\n# # self.final_layer = nn.ModuleList(self.final_layer)\n\n# def _make_layer(self, block, planes, blocks, stride=1):\n# downsample = None\n# if stride != 1 or self.inplanes != planes * block.expansion:\n# downsample = nn.Sequential(nn.Conv2d(self.inplanes, planes * block.expansion,\n# kernel_size=1, stride=stride, bias=False),\n# nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM))\n\n# layers = []\n# layers.append(block(self.inplanes, planes, stride, downsample))\n# self.inplanes = planes * block.expansion\n# for i in range(1, blocks):\n# layers.append(block(self.inplanes, planes))\n# return nn.Sequential(*layers)\n\n# def _get_deconv_cfg(self, deconv_kernel, index):\n# if deconv_kernel == 4:\n# padding = 1\n# output_padding = 0\n# elif deconv_kernel == 3:\n# padding = 1\n# output_padding = 1\n# elif deconv_kernel == 2:\n# padding = 0\n# output_padding = 0\n\n# return deconv_kernel, padding, output_padding\n\n# def _make_deconv_layer(self, num_layers, num_filters, num_kernels, inplanes):\n# assert num_layers == len(num_filters), \\\n# 'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n# assert num_layers == len(num_kernels), \\\n# 'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n \n# if self.verbose:\n# print(\"_make_deconv_layer inplanes: {}\".format(inplanes))\n\n# layers = []\n \n# for i in range(num_layers):\n# kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i], i)\n\n# planes = num_filters[i]\n# layers.append(nn.ConvTranspose2d(in_channels=inplanes,\n# out_channels=planes,\n# kernel_size=kernel,\n# stride=2,\n# padding=padding,\n# output_padding=output_padding,\n# bias=self.deconv_with_bias))\n# layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))\n# layers.append(nn.ReLU(inplace=True))\n# inplanes = planes\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# if self.verbose:\n# print(\"x shape: {}\".format(x.size()))\n\n# x = self.layer1(x)\n# if self.verbose:\n# print(\"x shape: {}\".format(x.size()))\n# x64 = self.layer2(x)\n# if self.verbose:\n# print(\"x64 shape: {}\".format(x64.size()))\n# x32 = self.layer3(x64)\n# if self.verbose:\n# print(\"x32 shape: {}\".format(x32.size()))\n# x16 = self.layer4(x32)\n# if self.verbose:\n# print(\"x16 shape: {}\".format(x16.size()))\n# # conbine deconv x from diff layer\n# x16 = self.deconv_layer16(x16)\n# x32 = self.deconv_layer32(x32)\n# x64 = self.deconv_layer64(x64)\n \n# x = torch.cat([x16,x32,x64],1)\n \n# if self.verbose:\n# print(\"x shape: {}\".format(x.size()))\n# out = [[self.hmap(x), self.regs(x), self.w_h_(x)]]\n# return out\n\n# def init_weights(self, num_layers):\n# for m in self.deconv_layer16.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n \n# for m in self.deconv_layer32.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n# for m in self.deconv_layer64.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n \n# for m in self.hmap.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.constant_(m.bias, -2.19)\n# for m in self.regs.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.normal_(m.weight, std=0.001)\n# nn.init.constant_(m.bias, 0)\n# for m in self.w_h_.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.normal_(m.weight, std=0.001)\n# nn.init.constant_(m.bias, 0)\n# url = model_urls['resnet{}'.format(num_layers)]\n# pretrained_state_dict = model_zoo.load_url(url)\n# print('=> loading pretrained model {}'.format(url))\n# self.load_state_dict(pretrained_state_dict, strict=False)\n\n\n# resnet_spec = {18: (BasicBlock, [2, 2, 2, 2]),\n# 34: (BasicBlock, [3, 4, 6, 3]),\n# 50: (Bottleneck, [3, 4, 6, 3]),\n# 101: (Bottleneck, [3, 4, 23, 3]),\n# 152: (Bottleneck, [3, 8, 36, 3])}\n\n# def resnet_18():\n# model = PoseResNet(BasicBlock, [2, 2, 2, 2], head_conv=64, num_classes=80,verbose=False)\n# model.init_weights(18)\n# return model\n\n# def get_pose_net(num_layers, head_conv, num_classes=80):\n# block_class, layers = resnet_spec[num_layers]\n\n# model = PoseResNet(block_class, layers, head_conv=head_conv, num_classes=num_classes, verbose= False)\n# # model.init_weights(num_layers)\n# return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# import torch\n# import torch.nn as nn\n# import torch.utils.model_zoo as model_zoo\n\n# BN_MOMENTUM = 0.1\n\n# model_urls = {'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n# 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n# 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n# 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n# 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', }\n\n\n# def conv3x3(in_planes, out_planes, stride=1):\n# \"\"\"3x3 convolution with padding\"\"\"\n# return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n\n\n# class BasicBlock(nn.Module):\n# expansion = 1\n\n# def __init__(self, inplanes, planes, stride=1, downsample=None):\n# super(BasicBlock, self).__init__()\n# self.conv1 = conv3x3(inplanes, planes, stride)\n# self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.relu = nn.ReLU(inplace=True)\n# self.conv2 = conv3x3(planes, planes)\n# self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.downsample = downsample\n# self.stride = stride\n\n# def forward(self, x):\n# residual = 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\n# if self.downsample is not None:\n# residual = self.downsample(x)\n\n# out += residual\n# out = self.relu(out)\n\n# return out\n\n\n# class Bottleneck(nn.Module):\n# expansion = 4\n\n# def __init__(self, inplanes, planes, stride=1, downsample=None):\n# super(Bottleneck, self).__init__()\n# self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n# self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n# self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n# self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n# self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM)\n# self.relu = nn.ReLU(inplace=True)\n# self.downsample = downsample\n# self.stride = stride\n\n# def forward(self, x):\n# residual = 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# if self.downsample is not None:\n# residual = self.downsample(x)\n\n# out += residual\n# out = self.relu(out)\n\n# return out\n\n\n# class PoseResNet(nn.Module):\n# def __init__(self, block, layers, head_conv, num_classes, verbose = True):\n# super(PoseResNet, self).__init__()\n# self.inplanes = 64\n# self.deconv_with_bias = False\n# self.num_classes = num_classes\n# self.verbose = verbose\n\n# self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n# self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\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, layers[0])\n# self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n# self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n# self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n\n# # used for deconv layers\n# self.deconv_layer16 = self._make_deconv_layer(1, [1024], [4],2048)\n# self.deconv_layer32 = self._make_deconv_layer(1, [512], [4],2048)\n# self.deconv_layer64 = self._make_deconv_layer(1, [256], [4],1024)\n# # self.final_layer = []\n\n# if head_conv > 0:\n# # heatmap layers\n# self.hmap = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, num_classes, kernel_size=1))\n# # self.hmap[-1].bias.data.fill_(-2.19)\n# # regression layers\n# self.regs = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, 2, kernel_size=1))\n# self.w_h_ = nn.Sequential(nn.Conv2d(256, head_conv, kernel_size=3, padding=1),\n# nn.ReLU(inplace=True),\n# nn.Conv2d(head_conv, 2, kernel_size=1))\n# else:\n# # heatmap layers\n# self.hmap = nn.Conv2d(in_channels=256, out_channels=num_classes, kernel_size=1)\n# # regression layers\n# self.regs = nn.Conv2d(in_channels=256, out_channels=2, kernel_size=1)\n# self.w_h_ = nn.Conv2d(in_channels=256, out_channels=2, kernel_size=1)\n\n# # self.final_layer = nn.ModuleList(self.final_layer)\n\n# def _make_layer(self, block, planes, blocks, stride=1):\n# downsample = None\n# if stride != 1 or self.inplanes != planes * block.expansion:\n# downsample = nn.Sequential(nn.Conv2d(self.inplanes, planes * block.expansion,\n# kernel_size=1, stride=stride, bias=False),\n# nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM))\n\n# layers = []\n# layers.append(block(self.inplanes, planes, stride, downsample))\n# self.inplanes = planes * block.expansion\n# for i in range(1, blocks):\n# layers.append(block(self.inplanes, planes))\n# return nn.Sequential(*layers)\n\n# def _get_deconv_cfg(self, deconv_kernel, index):\n# if deconv_kernel == 4:\n# padding = 1\n# output_padding = 0\n# elif deconv_kernel == 3:\n# padding = 1\n# output_padding = 1\n# elif deconv_kernel == 2:\n# padding = 0\n# output_padding = 0\n\n# return deconv_kernel, padding, output_padding\n\n# def _make_deconv_layer(self, num_layers, num_filters, num_kernels, inplanes):\n# assert num_layers == len(num_filters), \\\n# 'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n# assert num_layers == len(num_kernels), \\\n# 'ERROR: num_deconv_layers is different len(num_deconv_filters)'\n \n# if self.verbose:\n# print(\"_make_deconv_layer inplanes: {}\".format(inplanes))\n\n# layers = []\n \n# for i in range(num_layers):\n# kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i], i)\n\n# planes = num_filters[i]\n# layers.append(nn.ConvTranspose2d(in_channels=inplanes,\n# out_channels=planes,\n# kernel_size=kernel,\n# stride=2,\n# padding=padding,\n# output_padding=output_padding,\n# bias=self.deconv_with_bias))\n# layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))\n# layers.append(nn.ReLU(inplace=True))\n# inplanes = planes\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# if self.verbose:\n# print(\"x shape: {}\".format(x.size()))\n\n# x = self.layer1(x)\n# if self.verbose:\n# print(\"x shape: {}\".format(x.size()))\n# x64 = self.layer2(x)\n# if self.verbose:\n# print(\"x64 shape: {}\".format(x64.size()))\n# x32 = self.layer3(x64)\n# if self.verbose:\n# print(\"x32 shape: {}\".format(x32.size()))\n# x16 = self.layer4(x32)\n# if self.verbose:\n# print(\"x16 shape: {}\".format(x16.size()))\n# # conbine deconv x from diff layer\n# # x16 = self.deconv_layer16(x16)\n# # x32 = self.deconv_layer32(x32)\n# # x64 = self.deconv_layer64(x64)\n# up1 = self.deconv_layer16(x16)\n# if self.verbose:\n# print(\"up1 shape: {}\".format(up1.size()))\n# up1 = torch.cat([up1,x32],1)\n# if self.verbose:\n# print(\"up1 shape: {}\".format(up1.size()))\n# up2 = self.deconv_layer32(up1)\n# if self.verbose:\n# print(\"up2 shape: {}\".format(up2.size()))\n# up2 = torch.cat([up2,x64],1)\n# if self.verbose:\n# print(\"up2 shape: {}\".format(up2.size()))\n# x = self.deconv_layer64(up2)\n \n# # x = torch.cat([x16,x32,x64],1)\n \n# if self.verbose:\n# print(\"x shape: {}\".format(x.size()))\n# out = [[self.hmap(x), self.regs(x), self.w_h_(x)]]\n# return out\n\n# def init_weights(self, num_layers):\n# for m in self.deconv_layer16.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n \n# for m in self.deconv_layer32.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n# for m in self.deconv_layer64.modules():\n# if isinstance(m, nn.ConvTranspose2d):\n# nn.init.normal_(m.weight, std=0.001)\n# elif isinstance(m, nn.BatchNorm2d):\n# nn.init.constant_(m.weight, 1)\n# nn.init.constant_(m.bias, 0)\n \n# for m in self.hmap.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.constant_(m.bias, -2.19)\n# for m in self.regs.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.normal_(m.weight, std=0.001)\n# nn.init.constant_(m.bias, 0)\n# for m in self.w_h_.modules():\n# if isinstance(m, nn.Conv2d):\n# nn.init.normal_(m.weight, std=0.001)\n# nn.init.constant_(m.bias, 0)\n# url = model_urls['resnet{}'.format(num_layers)]\n# pretrained_state_dict = model_zoo.load_url(url)\n# print('=> loading pretrained model {}'.format(url))\n# self.load_state_dict(pretrained_state_dict, strict=False)\n\n\n# resnet_spec = {18: (BasicBlock, [2, 2, 2, 2]),\n# 34: (BasicBlock, [3, 4, 6, 3]),\n# 50: (Bottleneck, [3, 4, 6, 3]),\n# 101: (Bottleneck, [3, 4, 23, 3]),\n# 152: (Bottleneck, [3, 8, 36, 3])}\n\n# def resnet_18():\n# model = PoseResNet(BasicBlock, [2, 2, 2, 2], head_conv=64, num_classes=80,verbose=False)\n# model.init_weights(18)\n# return model\n\n# def get_pose_net(num_layers, head_conv, num_classes=80,verbose = False):\n# block_class, layers = resnet_spec[num_layers]\n\n# model = PoseResNet(block_class, layers, head_conv=head_conv, num_classes=num_classes, verbose= verbose)\n# # model.init_weights(num_layers)\n# return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## model2 hourglass","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport torch\nimport torch.nn as nn\n\n\nclass convolution(nn.Module):\n def __init__(self, k, inp_dim, out_dim, stride=1, with_bn=True):\n super(convolution, self).__init__()\n pad = (k - 1) // 2\n self.conv = nn.Conv2d(inp_dim, out_dim, (k, k), padding=(pad, pad), stride=(stride, stride), bias=not with_bn)\n self.bn = nn.BatchNorm2d(out_dim) if with_bn else nn.Sequential()\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n conv = self.conv(x)\n bn = self.bn(conv)\n relu = self.relu(bn)\n return relu\n\n\nclass residual(nn.Module):\n def __init__(self, k, inp_dim, out_dim, stride=1, with_bn=True):\n super(residual, self).__init__()\n\n self.conv1 = nn.Conv2d(inp_dim, out_dim, (3, 3), padding=(1, 1), stride=(stride, stride), bias=False)\n self.bn1 = nn.BatchNorm2d(out_dim)\n self.relu1 = nn.ReLU(inplace=True)\n\n self.conv2 = nn.Conv2d(out_dim, out_dim, (3, 3), padding=(1, 1), bias=False)\n self.bn2 = nn.BatchNorm2d(out_dim)\n\n self.skip = nn.Sequential(nn.Conv2d(inp_dim, out_dim, (1, 1), stride=(stride, stride), bias=False),\n nn.BatchNorm2d(out_dim)) \\\n if stride != 1 or inp_dim != out_dim else nn.Sequential()\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n conv1 = self.conv1(x)\n bn1 = self.bn1(conv1)\n relu1 = self.relu1(bn1)\n\n conv2 = self.conv2(relu1)\n bn2 = self.bn2(conv2)\n\n skip = self.skip(x)\n return self.relu(bn2 + skip)\n\n\n# inp_dim -> out_dim -> ... -> out_dim\ndef make_layer(kernel_size, inp_dim, out_dim, modules, layer, stride=1):\n layers = [layer(kernel_size, inp_dim, out_dim, stride=stride)]\n layers += [layer(kernel_size, out_dim, out_dim) for _ in range(modules - 1)]\n return nn.Sequential(*layers)\n\n\n# inp_dim -> inp_dim -> ... -> inp_dim -> out_dim\ndef make_layer_revr(kernel_size, inp_dim, out_dim, modules, layer):\n layers = [layer(kernel_size, inp_dim, inp_dim) for _ in range(modules - 1)]\n layers.append(layer(kernel_size, inp_dim, out_dim))\n return nn.Sequential(*layers)\n\n\n# key point layer\ndef make_kp_layer(cnv_dim, curr_dim, out_dim):\n return nn.Sequential(convolution(3, cnv_dim, curr_dim, with_bn=False),\n nn.Conv2d(curr_dim, out_dim, (1, 1)))\n\n\nclass kp_module(nn.Module):\n def __init__(self, n, dims, modules):\n super(kp_module, self).__init__()\n\n self.n = n\n\n curr_modules = modules[0]\n next_modules = modules[1]\n\n curr_dim = dims[0]\n next_dim = dims[1]\n\n # curr_mod x residual，curr_dim -> curr_dim -> ... -> curr_dim\n self.top = make_layer(3, curr_dim, curr_dim, curr_modules, layer=residual)\n self.down = nn.Sequential()\n # curr_mod x residual，curr_dim -> next_dim -> ... -> next_dim\n self.low1 = make_layer(3, curr_dim, next_dim, curr_modules, layer=residual, stride=2)\n # next_mod x residual，next_dim -> next_dim -> ... -> next_dim\n if self.n > 1:\n self.low2 = kp_module(n - 1, dims[1:], modules[1:])\n else:\n self.low2 = make_layer(3, next_dim, next_dim, next_modules, layer=residual)\n # curr_mod x residual，next_dim -> next_dim -> ... -> next_dim -> curr_dim\n self.low3 = make_layer_revr(3, next_dim, curr_dim, curr_modules, layer=residual)\n self.up = nn.Upsample(scale_factor=2)\n\n def forward(self, x):\n up1 = self.top(x)\n down = self.down(x)\n low1 = self.low1(down)\n low2 = self.low2(low1)\n low3 = self.low3(low2)\n up2 = self.up(low3)\n return up1 + up2\n\n\nclass exkp(nn.Module):\n def __init__(self, n, nstack, dims, modules, cnv_dim=256, num_classes=1):\n super(exkp, self).__init__()\n\n self.nstack = nstack\n self.num_classes = num_classes\n\n curr_dim = dims[0]\n\n self.pre = nn.Sequential(convolution(7, 3, 128, stride=2),\n residual(3, 128, curr_dim, stride=2))\n\n self.kps = nn.ModuleList([kp_module(n, dims, modules) for _ in range(nstack)])\n\n self.cnvs = nn.ModuleList([convolution(3, curr_dim, cnv_dim) for _ in range(nstack)])\n\n self.inters = nn.ModuleList([residual(3, curr_dim, curr_dim) for _ in range(nstack - 1)])\n\n self.inters_ = nn.ModuleList([nn.Sequential(nn.Conv2d(curr_dim, curr_dim, (1, 1), bias=False),\n nn.BatchNorm2d(curr_dim))\n for _ in range(nstack - 1)])\n self.cnvs_ = nn.ModuleList([nn.Sequential(nn.Conv2d(cnv_dim, curr_dim, (1, 1), bias=False),\n nn.BatchNorm2d(curr_dim))\n for _ in range(nstack - 1)])\n # heatmap layers\n self.hmap = nn.ModuleList([make_kp_layer(cnv_dim, curr_dim, num_classes) for _ in range(nstack)])\n for hmap in self.hmap:\n hmap[-1].bias.data.fill_(-2.19)\n\n # regression layers\n self.regs = nn.ModuleList([make_kp_layer(cnv_dim, curr_dim, 2) for _ in range(nstack)])\n self.w_h_ = nn.ModuleList([make_kp_layer(cnv_dim, curr_dim, 2) for _ in range(nstack)])\n\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, image):\n inter = self.pre(image)\n\n outs = []\n for ind in range(self.nstack):\n kp = self.kps[ind](inter)\n cnv = self.cnvs[ind](kp)\n\n if self.training or ind == self.nstack - 1:\n outs.append([self.hmap[ind](cnv), self.regs[ind](cnv), self.w_h_[ind](cnv)])\n\n if ind < self.nstack - 1:\n inter = self.inters_[ind](inter) + self.cnvs_[ind](cnv)\n inter = self.relu(inter)\n inter = self.inters[ind](inter)\n return outs\n\n\nget_hourglass = \\\n {'large_hourglass':\n exkp(n=5, nstack=2, dims=[256, 256, 384, 384, 384, 512], modules=[2, 2, 2, 2, 2, 4]),\n 'small_hourglass':\n exkp(n=5, nstack=1, dims=[256, 256, 384, 384, 384, 512], modules=[2, 2, 2, 2, 2, 4])}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define loss func ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef _neg_loss_slow(preds, targets):\n pos_inds = targets == 1 # todo targets > 1-epsilon ?\n neg_inds = targets < 1 # todo targets < 1-epsilon ?\n\n neg_weights = torch.pow(1 - targets[neg_inds], 4)\n\n loss = 0\n for pred in preds:\n pred = torch.clamp(torch.sigmoid(pred), min=1e-4, max=1 - 1e-4)\n pos_pred = pred[pos_inds]\n neg_pred = pred[neg_inds]\n\n pos_loss = torch.log(pos_pred) * torch.pow(1 - pos_pred, 2)\n neg_loss = torch.log(1 - neg_pred) * torch.pow(neg_pred, 2) * neg_weights\n\n num_pos = pos_inds.float().sum()\n pos_loss = pos_loss.sum()\n neg_loss = neg_loss.sum()\n\n if pos_pred.nelement() == 0:\n loss = loss - neg_loss\n else:\n loss = loss - (pos_loss + neg_loss) / num_pos\n return loss\n\n\ndef _neg_loss(preds, targets):\n ''' Modified focal loss. Exactly the same as CornerNet.\n Runs faster and costs a little bit more memory\n Arguments:\n preds (B x c x h x w)\n gt_regr (B x c x h x w)\n '''\n pos_inds = targets.eq(1).float()\n neg_inds = targets.lt(1).float()\n\n neg_weights = torch.pow(1 - targets, 4)\n\n loss = 0\n for pred in preds:\n pred = torch.clamp(torch.sigmoid(pred), min=1e-4, max=1 - 1e-4)\n pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds\n neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds\n\n num_pos = pos_inds.float().sum()\n pos_loss = pos_loss.sum()\n neg_loss = neg_loss.sum()\n\n if num_pos == 0:\n loss = loss - neg_loss\n else:\n loss = loss - (pos_loss + neg_loss) / num_pos\n return loss / len(preds)\n\n\ndef _reg_loss(regs, gt_regs, mask):\n mask = mask[:, :, None].expand_as(gt_regs).float()\n loss = sum(F.l1_loss(r * mask, gt_regs * mask, reduction='sum') / (mask.sum() + 1e-4) for r in regs)\n return loss / len(regs)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Add some utils","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import OrderedDict\n\ndef _gather_feature(feat, ind, mask=None):\n dim = feat.size(2)\n ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)\n feat = feat.gather(1, ind)\n if mask is not None:\n mask = mask.unsqueeze(2).expand_as(feat)\n feat = feat[mask]\n feat = feat.view(-1, dim)\n return feat\n\n\ndef _tranpose_and_gather_feature(feat, ind):\n feat = feat.permute(0, 2, 3, 1).contiguous()\n feat = feat.view(feat.size(0), -1, feat.size(3))\n feat = _gather_feature(feat, ind)\n return feat\n\n\ndef flip_tensor(x):\n return torch.flip(x, [3])\n\n\ndef _nms(heat, kernel=3):\n hmax = F.max_pool2d(heat, kernel, stride=1, padding=(kernel - 1) // 2)\n keep = (hmax == heat).float()\n return heat * keep\n\n\ndef _topk(scores, K=40):\n batch, cat, height, width = scores.size()\n\n topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)\n\n topk_inds = topk_inds % (height * width)\n topk_ys = (topk_inds / width).int().float()\n topk_xs = (topk_inds % width).int().float()\n\n topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)\n topk_clses = (topk_ind / K).int()\n topk_inds = _gather_feature(topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)\n topk_ys = _gather_feature(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)\n topk_xs = _gather_feature(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)\n\n return topk_score, topk_inds, topk_clses, topk_ys, topk_xs\n\n\ndef ctdet_decode(hmap, regs, w_h_, K=100):\n batch, cat, height, width = hmap.shape\n hmap=torch.sigmoid(hmap)\n\n # if flip test\n if batch > 1:\n hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2\n w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2\n regs = regs[0:1]\n\n batch = 1\n\n hmap = _nms(hmap) # perform nms on heatmaps\n\n scores, inds, clses, ys, xs = _topk(hmap, K=K)\n\n regs = _tranpose_and_gather_feature(regs, inds)\n regs = regs.view(batch, K, 2)\n xs = xs.view(batch, K, 1) + regs[:, :, 0:1]\n ys = ys.view(batch, K, 1) + regs[:, :, 1:2]\n\n w_h_ = _tranpose_and_gather_feature(w_h_, inds)\n w_h_ = w_h_.view(batch, K, 2)\n\n clses = clses.view(batch, K, 1).float()\n scores = scores.view(batch, K, 1)\n bboxes = torch.cat([xs - w_h_[..., 0:1] / 2,\n ys - w_h_[..., 1:2] / 2,\n xs + w_h_[..., 0:1] / 2,\n ys + w_h_[..., 1:2] / 2], dim=2)\n detections = torch.cat([bboxes, scores, clses], dim=2)\n return detections\n\ndef transform_preds(coords, center, scale, output_size):\n target_coords = np.zeros(coords.shape)\n trans = get_affine_transform(center, scale, 0, output_size, inv=1)\n for p in range(coords.shape[0]):\n target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)\n return target_coords","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def calculate_iou(gt, pr, form='pascal_voc') -> float:\n \"\"\"Calculates the Intersection over Union.\n\n Args:\n gt: (np.ndarray[Union[int, float]]) coordinates of the ground-truth box\n pr: (np.ndarray[Union[int, float]]) coordinates of the prdected box\n form: (str) gt/pred coordinates format\n - pascal_voc: [xmin, ymin, xmax, ymax]\n - coco: [xmin, ymin, w, h]\n Returns:\n (float) Intersection over union (0.0 <= iou <= 1.0)\n \"\"\"\n if form == 'coco':\n gt = gt.copy()\n pr = pr.copy()\n\n gt[2] = gt[0] + gt[2]\n gt[3] = gt[1] + gt[3]\n pr[2] = pr[0] + pr[2]\n pr[3] = pr[1] + pr[3]\n\n # Calculate overlap area\n dx = min(gt[2], pr[2]) - max(gt[0], pr[0]) + 1\n\n if dx < 0:\n return 0.0\n\n dy = min(gt[3], pr[3]) - max(gt[1], pr[1]) + 1\n\n if dy < 0:\n return 0.0\n\n overlap_area = dx * dy\n\n # Calculate union area\n union_area = (\n (gt[2] - gt[0] + 1) * (gt[3] - gt[1] + 1) +\n (pr[2] - pr[0] + 1) * (pr[3] - pr[1] + 1) -\n overlap_area\n )\n\n return overlap_area / union_area\n\n\ndef find_best_match(gts, pred, pred_idx, threshold=0.5, form='pascal_voc', ious=None) -> int:\n \"\"\"Returns the index of the 'best match' between the\n ground-truth boxes and the prediction. The 'best match'\n is the highest IoU. (0.0 IoUs are ignored).\n\n Args:\n gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n pred: (List[Union[int, float]]) Coordinates of the predicted box\n pred_idx: (int) Index of the current predicted box\n threshold: (float) Threshold\n form: (str) Format of the coordinates\n ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n Return:\n (int) Index of the best match GT box (-1 if no match above threshold)\n \"\"\"\n best_match_iou = -np.inf\n best_match_idx = -1\n\n for gt_idx in range(len(gts)):\n\n if gts[gt_idx][0] < 0:\n # Already matched GT-box\n continue\n\n iou = -1 if ious is None else ious[gt_idx][pred_idx]\n\n if iou < 0:\n iou = calculate_iou(gts[gt_idx], pred, form=form)\n\n if ious is not None:\n ious[gt_idx][pred_idx] = iou\n\n if iou < threshold:\n continue\n\n if iou > best_match_iou:\n best_match_iou = iou\n best_match_idx = gt_idx\n\n return best_match_idx\n\ndef calculate_precision(gts, preds, threshold=0.5, form='coco', ious=None) -> float:\n \"\"\"Calculates precision for GT - prediction pairs at one threshold.\n\n Args:\n gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n sorted by confidence value (descending)\n threshold: (float) Threshold\n form: (str) Format of the coordinates\n ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n Return:\n (float) Precision\n \"\"\"\n n = len(preds)\n tp = 0\n fp = 0\n\n # for pred_idx, pred in enumerate(preds_sorted):\n for pred_idx in range(n):\n\n best_match_gt_idx = find_best_match(gts, preds[pred_idx], pred_idx,\n threshold=threshold, form=form, ious=ious)\n\n if best_match_gt_idx >= 0:\n # True positive: The predicted box matches a gt box with an IoU above the threshold.\n tp += 1\n # Remove the matched GT box\n gts[best_match_gt_idx] = -1\n\n else:\n # No match\n # False positive: indicates a predicted box had no associated gt box.\n fp += 1\n\n # False negative: indicates a gt box had no associated predicted box.\n fn = (gts.sum(axis=1) > 0).sum()\n\n return tp / (tp + fp + fn)\n\n\ndef calculate_image_precision(gts, preds, thresholds=(0.5,), form='pascal_voc') -> float:\n \"\"\"Calculates image precision.\n\n Args:\n gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n sorted by confidence value (descending)\n thresholds: (float) Different thresholds\n form: (str) Format of the coordinates\n\n Return:\n (float) Precision\n \"\"\"\n n_threshold = len(thresholds)\n image_precision = 0.0\n\n ious = np.ones((len(gts), len(preds))) * -1\n # ious = None\n\n for threshold in thresholds:\n precision_at_threshold = calculate_precision(gts.copy(), preds, threshold=threshold,\n form=form, ious=ious)\n image_precision += precision_at_threshold / n_threshold\n\n return image_precision\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## eval and submit","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n# model = get_pose_net(num_layers=50, head_conv=64, num_classes=1)\n# model.to(device)\nmodel = get_hourglass['small_hourglass']\nmodel.to(device)\n# state_path = '../input/0625chkpt/chkpt-resnet50fpn-0626-1.pt'\n# state_path = \"../input/0625chkpt/chkpt-resnet50-0612-1.pt\"\nstate_path = \"../input/0625chkpt/chkpt-smallhourglass-0627-1.pt\"\nstate = torch.load(state_path)\nfor k in state:\n print(k)\nmodel.load_state_dict(state['net'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport os\nimport matplotlib.pyplot as plt\n\nDIR_INPUT = '/kaggle/input'\nDIR_TRAIN = f'{DIR_INPUT}/global-wheat-detection/train'\nDIR_TEST = f'{DIR_INPUT}/global-wheat-detection/test'\n\nCOCO_MEAN = [0.40789654, 0.44719302, 0.47026115]\nCOCO_STD = [0.28863828, 0.27408164, 0.27809835]\nCOCO_EIGEN_VALUES = [0.2141788, 0.01817699, 0.00341571]\nCOCO_EIGEN_VECTORS = [[-0.58752847, -0.69563484, 0.41340352],\n [-0.5832747, 0.00994535, -0.81221408],\n [-0.56089297, 0.71832671, 0.41158938]]\n\nclass WheatTest(torch.utils.data.Dataset):\n '''\n return [img, hmap, _w_h, regs, indx, ind_mask, center, scale, img_id]\n '''\n\n\n def __init__(self, dataframe, data_dir, fix_size=512):\n super(WheatTest, self).__init__()\n self.num_classes = 1\n self.data_dir = data_dir\n self.fix_size = fix_size\n\n self.data_rng = np.random.RandomState(123)\n self.mean = np.array(COCO_MEAN, dtype=np.float32)[None, None, :]\n self.std = np.array(COCO_STD, dtype=np.float32)[None, None, :]\n\n self.df = dataframe\n self.ids = dataframe['image_id'].unique()\n\n self.max_objs = 128\n self.padding = 31 # 31 for resnet/resdcn\n self.down_ratio = 4\n self.img_size = {'h': fix_size, 'w': fix_size}\n self.fmap_size = {'h': fix_size // self.down_ratio, 'w': fix_size // self.down_ratio}\n self.rand_scales = np.arange(0.6, 1.4, 0.1)\n self.gaussian_iou = 0.7\n\n def _get_border(self, border, size):\n i = 1\n while size - border // i <= border // i:\n i *= 2\n return border // i\n\n def __getitem__(self, idx):\n img_id = self.ids[idx]\n\n img_path = os.path.join(self.data_dir, 'test', self.ids[idx] + '.jpg')\n img = cv2.imread(img_path)\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n # img = cv2.resize(img,(self.fix_size,self.fix_size)) # convert to fix_size, default by 512\n height, width = img.shape[0], img.shape[1]\n center = np.array([width / 2., height / 2.], dtype=np.float32) # center of image\n\n scale = max(height, width) * 1.0\n \n\n trans_img = get_affine_transform(center, scale, 0, [self.img_size['w'], self.img_size['h']])\n\n img = cv2.warpAffine(img, trans_img, (self.img_size['w'], self.img_size['h']))\n \n img = img.astype(np.float32) / 255.\n\n\n\n img -= self.mean\n img /= self.std\n img = img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]\n\n\n return {'image': img,'c': center, 's': scale, 'img_id': img_id}\n\n def __len__(self):\n return len(self.ids)\n\n\ntest_df = pd.read_csv('../input/global-wheat-detection/sample_submission.csv')\n\ntest_dataset = WheatTest(test_df,'../input/global-wheat-detection')\n\ntest_loader = torch.utils.data.DataLoader(\n test_dataset,\n batch_size=1,\n shuffle=False,\n num_workers=0,\n)\n\n\nclass Wheat(torch.utils.data.Dataset):\n '''\n return [img, hmap, _w_h, regs, indx, ind_mask, center, scale, img_id]\n '''\n\n\n def __init__(self, dataframe, data_dir, train=True, transform=None, fix_size=512):\n super(Wheat, self).__init__()\n self.num_classes = 1\n self.transform = transform\n self.data_dir = data_dir\n self.fix_size = fix_size\n\n self.data_rng = np.random.RandomState(123)\n self.eig_val = np.array(COCO_EIGEN_VALUES, dtype=np.float32)\n self.eig_vec = np.array(COCO_EIGEN_VECTORS, dtype=np.float32)\n self.mean = np.array(COCO_MEAN, dtype=np.float32)[None, None, :]\n self.std = np.array(COCO_STD, dtype=np.float32)[None, None, :]\n\n self.df = dataframe\n self.ids = dataframe['image_id'].unique()\n self.train = train\n\n self.max_objs = 128\n self.padding = 31 # 31 for resnet/resdcn\n self.down_ratio = 4\n self.img_size = {'h': fix_size, 'w': fix_size}\n self.fmap_size = {'h': fix_size // self.down_ratio, 'w': fix_size // self.down_ratio}\n self.rand_scales = np.arange(0.6, 1.4, 0.1)\n self.gaussian_iou = 0.7\n\n def _get_border(self, border, size):\n i = 1\n while size - border // i <= border // i:\n i *= 2\n return border // i\n\n def __getitem__(self, idx):\n img_id = self.ids[idx]\n\n img_path = os.path.join(self.data_dir, self.ids[idx] + '.jpg')\n img = cv2.imread(img_path)\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n # img = cv2.resize(img,(self.fix_size,self.fix_size)) # convert to fix_size, default by 512\n height, width = img.shape[0], img.shape[1]\n center = np.array([width / 2., height / 2.], dtype=np.float32) # center of image\n\n scale = max(height, width) * 1.0\n\n flipped = False\n if self.train:\n scale = scale * np.random.choice(self.rand_scales)\n w_border = self._get_border(128, img.shape[1])\n h_border = self._get_border(128, img.shape[0])\n center[0] = np.random.randint(low=w_border, high=width - w_border)\n center[1] = np.random.randint(low=h_border, high=height - h_border)\n\n if np.random.random() < 0.5:\n flipped = True\n img = img[:, ::-1, :]\n center[0] = width - center[0] - 1\n\n trans_img = get_affine_transform(center, scale, 0, [self.img_size['w'], self.img_size['h']])\n\n img = cv2.warpAffine(img, trans_img, (self.img_size['w'], self.img_size['h']))\n\n annos = self.df[self.df['image_id'].isin([self.ids[idx]])]\n\n bboxes = annos[['x', 'y', 'w', 'h']].values\n bboxes[:, 2:] += bboxes[:, :2] # xywh to xyxy\n gt = bboxes.copy()\n labels = np.zeros(len(bboxes)).astype(np.uint8)\n img = img.astype(np.float32) / 255.\n\n if self.train:\n color_aug(self.data_rng, img, self.eig_val, self.eig_vec)\n\n\n img -= self.mean\n img /= self.std\n img = img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]\n\n trans_fmap = get_affine_transform(center, scale, 0, [self.fmap_size['w'], self.fmap_size['h']])\n hmap = np.zeros((self.num_classes, self.fmap_size['h'], self.fmap_size['w']), dtype=np.float32) # heatmap\n w_h_ = np.zeros((self.max_objs, 2), dtype=np.float32) # width and height\n regs = np.zeros((self.max_objs, 2), dtype=np.float32) # regression\n inds = np.zeros((self.max_objs,), dtype=np.int64)\n ind_masks = np.zeros((self.max_objs,), dtype=np.uint8)\n\n # ###\n # fig = plt.figure()\n\n # ax1 = fig.add_subplot(121)\n # ax1.imshow(trans_img)\n\n # ax2 = fig.add_subplot(122)\n # ax2.imshow(trans_fmap)\n # ###\n for k, (bbox, label) in enumerate(zip(bboxes, labels)):\n if flipped:\n bbox[[0, 2]] = width - bbox[[2, 0]] - 1\n bbox[:2] = affine_transform(bbox[:2], trans_fmap)\n bbox[2:] = affine_transform(bbox[2:], trans_fmap)\n bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.fmap_size['w'] - 1)\n bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.fmap_size['h'] - 1)\n\n h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]\n\n if h > 0 and w > 0:\n obj_c = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)\n obj_c_int = obj_c.astype(np.int32)\n\n radius = max(0, int(gaussian_radius((np.ceil(h), np.ceil(w)), self.gaussian_iou)))\n draw_umich_gaussian(hmap[label], obj_c_int, radius)\n w_h_[k] = 1. * w, 1. * h\n regs[k] = obj_c - obj_c_int # discretization error\n inds[k] = obj_c_int[1] * self.fmap_size['w'] + obj_c_int[0]\n ind_masks[k] = 1\n # ###\n # fig = plt.figure()\n\n # ax1 = fig.add_subplot(121)\n # ax1.imshow(img.transpose(1,2,0))\n\n # ax2 = fig.add_subplot(122)\n # ax2.imshow(hmap.transpose(1,2,0).squeeze(2))\n # ###\n\n return {'image': img,\n 'hmap': hmap, 'w_h_': w_h_, 'regs': regs, 'inds': inds, 'ind_masks': ind_masks,\n 'c': center, 's': scale, 'img_id': img_id,'boxes':gt}\n\n def __len__(self):\n return len(self.ids)\n \n \ncsv_path = '/kaggle/input/global-wheat-detection/train.csv'\ndf = pd.read_csv(csv_path)\n\ndef process_bbox(df):\n df['bbox'] = df['bbox'].apply(lambda x: eval(x))\n df['x'] = df['bbox'].apply(lambda x: x[0])\n df['y'] = df['bbox'].apply(lambda x: x[1])\n df['w'] = df['bbox'].apply(lambda x: x[2])\n df['h'] = df['bbox'].apply(lambda x: x[3])\n df['x'] = df['x'].astype(np.float)\n df['y'] = df['y'].astype(np.float)\n df['w'] = df['w'].astype(np.float)\n df['h'] = df['h'].astype(np.float)\n\n df.drop(columns=['bbox'],inplace=True)\n# df.reset_index(drop=True)\n return df\n\n\ndf_new = process_bbox(df)\n\nimage_ids = df_new['image_id'].unique()\ntrain_ids = image_ids[0:int(0.8*len(image_ids))]\nval_ids = image_ids[int(0.8*len(image_ids)):]\ntrain_df = df_new[df_new['image_id'].isin(train_ids)]\nval_df = df_new[df_new['image_id'].isin(val_ids)]\n\nval_set = Wheat(val_df,DIR_TRAIN,False)\n\nval_loader = torch.utils.data.DataLoader(\n val_set,\n batch_size=1,\n shuffle=True, \n num_workers=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import time\nresults1 = []\nresults2 = []\nsocre_bound = 0.2\n\nt0 = time.perf_counter()\n\nfor sample in val_loader:\n for k in sample:\n if k != 'img_id' and k != 'boxes':\n sample[k] = sample[k].to(device)\n\n with torch.no_grad():\n img = sample['image']\n img_id = sample['img_id'][0]\n \n output = model(img)[-1]\n# print('using: {}ms'.format(1000*dur))\n \n dets = ctdet_decode(*output, K=50)\n dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])[0]\n dets[:, :2] = transform_preds(dets[:, 0:2],\n sample['c'].cpu().numpy(),\n sample['s'].cpu().numpy(),\n (128,128))\n dets[:, 2:4] = transform_preds(dets[:, 2:4],\n sample['c'].cpu().numpy(),\n sample['s'].cpu().numpy(),\n (128,128))\n# dur = time.perf_counter() - t0\n \n \n mean = np.array([0.40789654, 0.44719302, 0.47026115], dtype=np.float32).reshape(1, 1, 3)\n std = np.array([0.28863828, 0.27408164, 0.27809835], dtype=np.float32).reshape(1, 1, 3)\n boxes = dets[:,:4]\n scores = dets[:,-2]\n gt = sample['boxes'].numpy()[0]\n# print(gt.shape)\n\n# c_img = img.cpu().numpy()\n# c_img = c_img[0].transpose(1,2,0)\n# c_img = np.ascontiguousarray(c_img)\n \n# s_boxes = []\n# s_scores = []\n\n boxes = np.clip(boxes, 0 ,1024).astype(np.int32)\n# print(boxes)\n# print(gt)\n precision1 = calculate_image_precision(gt,boxes,(0.5,))\n results1.append(precision1)\n precision2 = calculate_image_precision(gt,boxes,(0.75,))\n results2.append(precision2)\n \n \n if precision1 >0.7:\n print(\"0.5mAP: {}\".format(np.mean(precision1)))\n print(\"0.75mAP: {}\".format(np.mean(precision2)))\n c_img = img.cpu().numpy()\n c_img = c_img[0].transpose(1,2,0)\n c_img = np.ascontiguousarray(c_img)\n c_img = (c_img*std)+mean\n c_img = cv2.resize(c_img,(1024,1024))\n for i, box in enumerate(boxes):\n box = box.astype(np.int32)\n# if dets[i-1][-2]<0.2:\n# continue\n cv2.rectangle(c_img,\n (box[0],box[1]),\n (box[2],box[3]),\n (200,200,200), 2)\n\n\n for i, box in enumerate(gt):\n box = box.astype(np.int32)\n# if dets[i-1][-2]<0.2:\n# continue\n cv2.rectangle(c_img,\n (int(box[0]),int(box[1])),\n (int(box[2]),int(box[3])),\n (200,0,0), 2)\n\n fig = plt.figure(figsize=(8,8))\n ax1 = fig.add_subplot(111)\n ax1.imshow(c_img)\n \nprint(\"TOTAL 0.5mAP: {}\".format(np.mean(results1)))\nprint(\"TOTAL 0.75mAP: {}\".format(np.mean(results2)))\n# result = {\n# 'image_id': img_id,\n# 'PredictionString': format_prediction_string(boxes, scores)\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'])\n# a = test_df.head()\n# test_df.to_csv('submission.csv', index=False)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}