{"cells":[{"metadata":{},"cell_type":"markdown","source":"# CenterNet starter codes in Pytorch!\n\n## Why CenterNet?\n\n![](https://github.com/xingyizhou/CenterNet/raw/master/readme/fig2.png)\n\n1) CenterNet is simple and easy to customize compared to faster rcnns.\n\n2) CenterNet is fast, since it's a single-stage detector.\n\n3) Can be accurate as the faster rcnn if trained properly!\n\n### This starterkit trains a simple CenterNet. \nThe network is underfitting, so you can extend this by:\n\na) training longer..\n\nb) adding augumentations..\n\nc) making the network deeper..\n\n### Please upvote if this notebook helps you, Thanks!!"},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm_notebook as tqdm\n\nimport torch\nimport torch.nn as nn\nimport torchvision","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Hyper-params\nMODEL_PATH = \"\"\ninput_size = 512\nIN_SCALE = 1024//input_size \nMODEL_SCALE = 4\nbatch_size = 2\nmodel_name = \"resnet18\"\nTRAIN = True # True for training","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare labels"},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = '../input/global-wheat-detection/'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\ntrain_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Split train-test\nfrom sklearn.model_selection import train_test_split\n# Split by unique image ids.\nimage_ids = train_df['image_id'].unique()\ntrain_id, test_id = train_test_split(image_ids, test_size=0.2, random_state=777)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Convert boxes to heatmap\nWHile the labels are given in COCO format, we need to convert them to heatmaps in order to train CenterNets."},{"metadata":{"trusted":true},"cell_type":"code","source":"# show image\nimg_id = train_id[0]\nimg = cv2.imread(os.path.join(DIR_INPUT,\"train\", img_id+\".jpg\"))\nimg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\nplt.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# get targets\ntarget = train_df[train_df['image_id']==img_id]\n# convert targets to its center.\ntry:\n    center = np.array([target[\"x\"]+target[\"w\"]//2, target[\"y\"]+target[\"h\"]//2]).T\nexcept:\n    center = np.array([int(target[\"x\"]+target[\"w\"]//2), int(target[\"y\"]+target[\"h\"]//2)]).T.reshape(1,2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's plot center points into the image."},{"metadata":{"trusted":true},"cell_type":"code","source":"# plot centers on image\nplt.figure(figsize=(14,14))\nplt.imshow(img)\nfor x in center:\n    plt.scatter(x[0], x[1], color='red', s=100)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Make heatmaps using the utility functions from the centernet repo\ndef draw_msra_gaussian(heatmap, center, sigma=2):\n  tmp_size = sigma * 6\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\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\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Wrapped heatmap function\ndef make_hm_regr(target):\n    # make output heatmap for single class\n    hm = np.zeros([input_size//MODEL_SCALE, input_size//MODEL_SCALE])\n    # make regr heatmap \n    regr = np.zeros([2, input_size//MODEL_SCALE, input_size//MODEL_SCALE])\n    \n    if len(target) == 0:\n        return hm, regr\n    \n    try:\n        center = np.array([target[\"x\"]+target[\"w\"]//2, target[\"y\"]+target[\"h\"]//2, \n                       target[\"w\"], target[\"h\"]\n                      ]).T\n    except:\n        center = np.array([int(target[\"x\"]+target[\"w\"]//2), int(target[\"y\"]+target[\"h\"]//2), \n                       int(target[\"w\"]), int(target[\"h\"])\n                      ]).T.reshape(1,4)\n    \n    # make a center point\n    # try gaussian points.\n    for c in center:\n        hm = draw_msra_gaussian(hm, [int(c[0])//MODEL_SCALE//IN_SCALE, int(c[1])//MODEL_SCALE//IN_SCALE], \n                                sigma=np.clip(c[2]*c[3]//2000, 2, 4))    \n\n    # convert targets to its center.\n    regrs = center[:, 2:]/input_size/IN_SCALE\n\n    # plot regr values to mask\n    for r, c in zip(regrs, center):\n        for i in range(-2, 3):\n            for j in range(-2, 3):\n                try:\n                    regr[:, int(c[0])//MODEL_SCALE//IN_SCALE+i, \n                         int(c[1])//MODEL_SCALE//IN_SCALE+j] = r\n                except:\n                    pass\n    regr[0] = regr[0].T; regr[1] = regr[1].T;\n    return hm, regr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def pred2box(hm, regr, thresh=0.99):\n    # make binding box from heatmaps\n    # thresh: threshold for logits.\n        \n    # get center\n    pred = hm > thresh\n    pred_center = np.where(hm>thresh)\n    # get regressions\n    pred_r = regr[:,pred].T\n\n    # wrap as boxes\n    # [xmin, ymin, width, height]\n    # size as original image.\n    boxes = []\n    scores = hm[pred]\n    for i, b in enumerate(pred_r):\n        arr = np.array([pred_center[1][i]*MODEL_SCALE-b[0]*input_size//2, pred_center[0][i]*MODEL_SCALE-b[1]*input_size//2, \n                      int(b[0]*input_size), int(b[1]*input_size)])\n        arr = np.clip(arr, 0, input_size)\n        # filter \n        #if arr[0]<0 or arr[1]<0 or arr[0]>input_size or arr[1]>input_size:\n            #pass\n        boxes.append(arr)\n    return np.asarray(boxes), scores","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# functions for plotting results\ndef showbox(img, hm, regr, thresh=0.9):\n    boxes, _ = pred2box(hm, regr, thresh=thresh)\n    print(\"preds:\",boxes.shape)\n    sample = img\n\n    for box in boxes:\n        # upper-left, lower-right\n        cv2.rectangle(sample,\n                      (int(box[0]), int(box[1]+box[3])),\n                      (int(box[0]+box[2]), int(box[1])),\n                      (220, 0, 0), 3)\n    return sample\n\ndef showgtbox(img, hm, regr, thresh=0.9):\n    boxes, _ = pred2box(hm, regr, thresh=thresh)\n    print(\"GT boxes:\", boxes.shape)\n    sample = img\n\n    for box in boxes:\n        cv2.rectangle(sample,\n                      (int(box[0]), int(box[1]+box[3])),\n                      (int(box[0]+box[2]), int(box[1])),\n                      (0, 220, 0), 3)\n    return sample","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## compile data processing\nLet's first make sure that we can convert targets to heatmaps and bring it back to boxes."},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"# show image\n#img_id = train_df[\"image_id\"][]\nimg_id = train_id[0]\nimg = cv2.imread(os.path.join(DIR_INPUT,\"train\", img_id+\".jpg\"))\nimg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\nimg = cv2.resize(img, (input_size, input_size))\nsample = img\n\n# get labels\ntarget = train_df[train_df['image_id']==img_id]\n\n# convert target to heatmaps\nhm, regr = make_hm_regr(target)\n\n# get boxes\nboxes, _ = pred2box(hm, regr)\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\nsample = showbox(sample, hm, regr, 0.99)\nplt.imshow(sample)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Make dataset\nThen, let's make a dataset and a dataloader for training."},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision import transforms\n\nclass Normalize(object):\n    def __init__(self):\n        self.mean=[0.485, 0.456, 0.406]\n        self.std=[0.229, 0.224, 0.225]\n        self.norm = transforms.Normalize(self.mean, self.std)\n    def __call__(self, image):\n        image = image.astype(np.float32)/255\n        axis = (0,1)\n        image -= self.mean\n        image /= self.std\n        return image\n    \n# pool duplicates\ndef pool(data):\n    stride = 3\n    for y in np.arange(1,data.shape[1]-1, stride):\n        for x in np.arange(1, data.shape[0]-1, stride):\n            a_2d = data[x-1:x+2, y-1:y+2]\n            max = np.asarray(np.unravel_index(np.argmax(a_2d), a_2d.shape))            \n            for c1 in range(3):\n                for c2 in range(3):\n                    #print(c1,c2)\n                    if not (c1== max[0] and c2 == max[1]):\n                        data[x+c1-1, y+c2-1] = -1\n    return data\n\nclass WheatDataset(torch.utils.data.Dataset):\n    def __init__(self, img_id, labels, transform=None):\n        self.img_id = img_id\n        self.labels = labels\n        if transform:\n            self.transform = transform\n        self.normalize = Normalize()\n        \n    def __len__(self):\n        return len(self.img_id)\n\n    def __getitem__(self, idx):\n        img = cv2.imread(os.path.join(DIR_INPUT,\"train\", self.img_id[idx]+\".jpg\"))\n        img = cv2.resize(img, (input_size, input_size))\n        img = self.normalize(img)\n        img = img.transpose([2,0,1])\n        target = self.labels[self.labels['image_id']==self.img_id[idx]]\n        hm, regr = make_hm_regr(target)\n        return img, hm, regr\n\n# Submission\nclass WheatDatasetTest(torch.utils.data.Dataset):\n    def __init__(self, image_dir, transform=None):\n        self.image_dir = image_dir\n        self.img_id = os.listdir(self.image_dir)\n        if transform:\n            self.transform = transform\n        self.normalize = Normalize()\n        \n    def __len__(self):\n        return len(self.img_id)\n\n    def __getitem__(self, idx):\n        img = cv2.imread(os.path.join(self.image_dir, self.img_id[idx]))\n        img = cv2.resize(img, (input_size, input_size))\n        img = self.normalize(img)\n        img = img.transpose([2,0,1])\n        return img, self.img_id[idx]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"traindataset = WheatDataset(train_id, train_df)\nvaldataset = WheatDataset(test_id, train_df)\ntestdataset = WheatDatasetTest('../input/global-wheat-detection/test')\n\n# Test dataset\nimg, hm, regr = traindataset[0]\nplt.imshow(img.transpose([1,2,0]))\nplt.show()\nimg.std()\nplt.imshow(hm)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Pack to dataloaders\ntrain_loader = torch.utils.data.DataLoader(traindataset,batch_size=batch_size,shuffle=True, num_workers=0)\nval_loader = torch.utils.data.DataLoader(valdataset,batch_size=batch_size,shuffle=True, num_workers=0)\ntest_loader = torch.utils.data.DataLoader(testdataset,batch_size=batch_size,shuffle=False, num_workers=0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Centernet model\nThe model is based on [PKU centernets](https://www.kaggle.com/hocop1/centernet-baseline/data)."},{"metadata":{"trusted":true},"cell_type":"code","source":"class double_conv(nn.Module):\n    '''(conv => BN => ReLU) * 2'''\n    def __init__(self, in_ch, out_ch):\n        super(double_conv, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        x = self.conv(x)\n        return x\n\nclass up(nn.Module):\n    def __init__(self, in_ch, out_ch, bilinear=True):\n        super(up, self).__init__()\n        if bilinear:\n            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n        else:\n            self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)\n        self.conv = double_conv(in_ch, out_ch)\n        \n    def forward(self, x1, x2=None):\n        x1 = self.up(x1)\n        if x2 is not None:\n            x = torch.cat([x2, x1], dim=1)\n            # input is CHW\n            diffY = x2.size()[2] - x1.size()[2]\n            diffX = x2.size()[3] - x1.size()[3]\n\n            x1 = F.pad(x1, (diffX // 2, diffX - diffX//2,\n                            diffY // 2, diffY - diffY//2))\n        else:\n            x = x1\n        x = self.conv(x)\n        return x\n\nclass centernet(nn.Module):\n    def __init__(self, n_classes=1, model_name=\"resnet18\"):\n        super(centernet, self).__init__()\n        # create backbone.\n        basemodel = torchvision.models.resnet18(pretrained=False) # turn this on for training\n        basemodel = nn.Sequential(*list(basemodel.children())[:-2])\n        # set basemodel\n        self.base_model = basemodel\n        \n        if model_name == \"resnet34\" or model_name==\"resnet18\":\n            num_ch = 512\n        else:\n            num_ch = 2048\n        \n        self.up1 = up(num_ch, 512)\n        self.up2 = up(512, 256)\n        self.up3 = up(256, 256)\n        # output classification\n        self.outc = nn.Conv2d(256, n_classes, 1)\n        # output residue\n        self.outr = nn.Conv2d(256, 2, 1)\n        \n    def forward(self, x):\n        batch_size = x.shape[0]\n        \n        x = self.base_model(x)\n        \n        # Add positional info        \n        x = self.up1(x)\n        x = self.up2(x)\n        x = self.up3(x)\n        outc = self.outc(x)\n        outr = self.outr(x)\n        return outc, outr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = centernet()\n# Check if it runs correctly\nmodel(torch.rand(1,3,512,512))[0].size()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define optimizers"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Gets the GPU if there is one, otherwise the cpu\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\n\n# Optimizer\nimport torch.optim as optim\noptimizer = optim.Adam(model.parameters(), lr=1e-4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"# From centernet repo\ndef neg_loss(pred, gt):\n  ''' Modified focal loss. Exactly the same as CornerNet.\n      Runs faster and costs a little bit more memory\n    Arguments:\n      pred (batch x c x h x w)\n      gt_regr (batch x c x h x w)\n  '''\n  pred = pred.unsqueeze(1).float()\n  gt = gt.unsqueeze(1).float()\n\n  pos_inds = gt.eq(1).float()\n  neg_inds = gt.lt(1).float()\n  neg_weights = torch.pow(1 - gt, 4)\n\n  loss = 0\n\n  pos_loss = torch.log(pred + 1e-12) * torch.pow(1 - pred, 3) * pos_inds\n  neg_loss = torch.log(1 - pred + 1e-12) * torch.pow(pred, 3) * 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\n\ndef _reg_loss(regr, gt_regr, mask):\n  ''' L1 regression loss\n    Arguments:\n      regr (batch x max_objects x dim)\n      gt_regr (batch x max_objects x dim)\n      mask (batch x max_objects)\n  '''\n  num = mask.float().sum()\n  #print(gt_regr.size())\n  mask = mask.sum(1).unsqueeze(1).expand_as(gt_regr)\n  #print(mask.size())\n\n  regr = regr * mask\n  gt_regr = gt_regr * mask\n    \n  regr_loss = nn.functional.smooth_l1_loss(regr, gt_regr, size_average=False)\n  regr_loss = regr_loss / (num + 1e-4)\n  return regr_loss\n  \ndef centerloss(prediction, mask, regr,weight=0.4, size_average=True):\n    # Binary mask loss\n    pred_mask = torch.sigmoid(prediction[:, 0])\n    mask_loss = neg_loss(pred_mask, mask)\n    \n    # Regression L1 loss\n    pred_regr = prediction[:, 1:]\n    regr_loss = (torch.abs(pred_regr - regr).sum(1) * mask).sum(1).sum(1) / mask.sum(1).sum(1)\n    regr_loss = regr_loss.mean(0)\n  \n    # Sum\n    loss = mask_loss +regr_loss\n    if not size_average:\n        loss *= prediction.shape[0]\n    return loss ,mask_loss , regr_loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train function"},{"metadata":{},"cell_type":"markdown","source":"![](https://i.imgflip.com/26d7zv.jpg)"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(epoch):\n    model.train()\n    print('epochs {}/{} '.format(epoch+1,epochs))\n    running_loss = 0.0\n    running_mask = 0.0\n    running_regr = 0.0\n    t = tqdm(train_loader)\n    rd = np.random.rand()\n    \n    for idx, (img, hm, regr) in enumerate(t):       \n        # send to gpu\n        img = img.to(device)\n        hm_gt = hm.to(device)\n        regr_gt = regr.to(device)\n        # set opt\n        optimizer.zero_grad()\n        \n        # run model\n        hm, regr = model(img)\n        preds = torch.cat((hm, regr), 1)\n            \n        loss, mask_loss, regr_loss = centerloss(preds, hm_gt, regr_gt)\n        # misc\n        running_loss += loss\n        running_mask += mask_loss\n        running_regr += regr_loss\n        \n        loss.backward()\n        optimizer.step()\n        \n        t.set_description(f't (l={running_loss/(idx+1):.3f})(m={running_mask/(idx+1):.4f})(r={running_regr/(idx+1):.4f})')\n        \n    #scheduler.step()\n    print('train loss : {:.4f}'.format(running_loss/len(train_loader)))\n    print('maskloss : {:.4f}'.format(running_mask/(len(train_loader))))\n    print('regrloss : {:.4f}'.format(running_regr/(len(train_loader))))\n    \n    # save logs\n    log_epoch = {'epoch': epoch+1, 'lr': optimizer.state_dict()['param_groups'][0]['lr'],\n                    'loss': running_loss/len(train_loader), \"mask\": running_mask/(len(train_loader)), \n                 \"regr\": running_regr/(len(train_loader))}\n    logs.append(log_epoch)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train model\nTo save time, only trained for 5 epochs"},{"metadata":{"trusted":true},"cell_type":"code","source":"#epochs=100\nepochs = 5\nlogs = []\nlogs_eval = []\n\nif TRAIN:\n    for epoch in range(epochs):\n        train(epoch)\nelse:\n    model.load_state_dict(torch.load(MODEL_PATH))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Show predictions"},{"metadata":{"trusted":true},"cell_type":"code","source":"for id in range(10):\n    img, hm_gt, regr_gt = valdataset[id]\n    img = torch.from_numpy(img)\n    with torch.no_grad():\n        hm, regr = model(img.to(device).float().unsqueeze(0))\n\n    \n    hm = hm.cpu().numpy().squeeze(0).squeeze(0)\n    regr = regr.cpu().numpy().squeeze(0)\n\n    # show image\n    img_id = test_id[id]\n    img = cv2.imread(os.path.join(DIR_INPUT,\"train\", img_id+\".jpg\"))\n    img = cv2.resize(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), (input_size, input_size))\n\n    # get boxes\n    hm = torch.sigmoid(torch.from_numpy(hm)).numpy()\n    hm = pool(hm)\n    plt.imshow(hm>0.6)\n    plt.show()\n    sample = showbox(img, hm, regr, 0.6)\n    \n    # show gt\n    sample = showgtbox(sample, hm_gt, regr_gt, 0.99)\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    plt.imshow(sample)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# For submissions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_prediction_string(boxes, scores):\n    pred_strings = []\n    for s, b in zip(scores, boxes.astype(int)):\n        # xmin, ymin, w, h\n        pred_strings.append(f'{s:.4f} {b[0]*IN_SCALE} {b[1]*IN_SCALE} {b[2]*IN_SCALE} {b[3]*IN_SCALE}')\n\n    return \" \".join(pred_strings)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"thresh = 0.6\nresults = []\n\nfor images, image_ids in tqdm(test_loader):\n\n    images = images.to(device)\n    with torch.no_grad():\n        hms, regrs = model(images)\n\n    for hm, regr, image_id in zip(hms, regrs, image_ids):\n        # process predictions\n        hm = hm.cpu().numpy().squeeze(0)\n        regr = regr.cpu().numpy()\n        hm = torch.sigmoid(torch.from_numpy(hm)).numpy()\n        hm = pool(hm)\n\n        boxes, scores = pred2box(hm, regr, thresh)\n\n        preds_sorted_idx = np.argsort(scores)[::-1]\n        boxes_sorted = boxes[preds_sorted_idx]\n        scores_sorted = scores[preds_sorted_idx]\n        \n        result = {\n            'image_id': image_id[:-4],\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'])\ntest_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"That's it!"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.5"}},"nbformat":4,"nbformat_minor":4}