{"cells":[{"metadata":{},"cell_type":"markdown","source":"Jeremy had [tweeted](https://twitter.com/jeremyphoward/status/1208135410733309952) earlier that FastAI v2 can be added to existing Pytorch Code, lets see how we can incorporate Centernet with the latest codebase."},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"#Install FastAI2 Library\n!pip install git+https://github.com/fastai/fastai2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Usual Imports\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom tqdm import tqdm\nfrom functools import reduce, partial\n\nimport os\nfrom sklearn.model_selection import train_test_split\nfrom scipy.optimize import minimize\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils import model_zoo\nimport random\nimport gc\nimport os.path\n\nfrom fastai2.basics import *\nfrom fastai2.callback.all import *\nfrom fastai2.vision.all import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#List of variables used in the code\n\n#Random seed setting\ntorch.manual_seed(42)\nrandom.seed(42) \n\n# Indicator to set all dataset is to be loaded\nfull=0 \n\n#Limit Training Datasize to 400\ntr_limit = 400\n\n#To track the loss\nbest_loss = torch.FloatTensor()\n\nn_epochs = 10\nBATCH_SIZE = 4\nSWITCH_LOSS_EPOCH = 5\nhistory = pd.DataFrame()\nDISTANCE_THRESH_CLEAR = 2\n\nPATH = '/kaggle/input/pku-autonomous-driving/'\nIMG_WIDTH = 2048 - 512\nIMG_HEIGHT = IMG_WIDTH // 16 * 5\nMODEL_SCALE = 8","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Code borrowed from other notebooks, credits in the references section\ntrain = pd.read_csv(PATH + 'train.csv')\ntest = pd.read_csv(PATH + 'sample_submission.csv')\n\ncamera_matrix = np.array([[2304.5479, 0,  1686.2379],\n                          [0, 2305.8757, 1354.9849],\n                          [0, 0, 1]], dtype=np.float32)\ncamera_matrix_inv = np.linalg.inv(camera_matrix)\n\n\ndef str2coords(s, names=['id', 'yaw', 'pitch', 'roll', 'x', 'y', 'z']):\n    '''\n    Input:\n        s: PredictionString (e.g. from train dataframe)\n        names: array of what to extract from the string\n    Output:\n        list of dicts with keys from `names`\n    '''\n    coords = []\n    for l in np.array(s.split()).reshape([-1, 7]):\n        coords.append(dict(zip(names, l.astype('float'))))\n        if 'id' in coords[-1]:\n            coords[-1]['id'] = int(coords[-1]['id'])\n    return coords\n\n\n#############\n\ndef imread(path, fast_mode=False):\n    img = cv2.imread(path)\n    if not fast_mode and img is not None and len(img.shape) == 3:\n        img = np.array(img[:, :, ::-1])\n    return img\n#############\n\n#############\n\n\n#############\ndef rotate(x, angle):\n    x = x + angle\n    x = x - (x + np.pi) // (2 * np.pi) * 2 * np.pi\n    return x\n\n#############\n\n\ndef get_img_coords(s):\n    '''\n    Input is a PredictionString (e.g. from train dataframe)\n    Output is two arrays:\n        xs: x coordinates in the image\n        ys: y coordinates in the image\n    '''\n    coords = str2coords(s)\n    xs = [c['x'] for c in coords]\n    ys = [c['y'] for c in coords]\n    zs = [c['z'] for c in coords]\n    P = np.array(list(zip(xs, ys, zs))).T\n    img_p = np.dot(camera_matrix, P).T\n    img_p[:, 0] /= img_p[:, 2]\n    img_p[:, 1] /= img_p[:, 2]\n    img_xs = img_p[:, 0]\n    img_ys = img_p[:, 1]\n    img_zs = img_p[:, 2] # z = Distance from the camera\n    return img_xs, img_ys\n\n\n#############\nfrom math import sin, cos\n\n# convert euler angle to rotation matrix\ndef euler_to_Rot(yaw, pitch, roll):\n    Y = np.array([[cos(yaw), 0, sin(yaw)],\n                  [0, 1, 0],\n                  [-sin(yaw), 0, cos(yaw)]])\n    P = np.array([[1, 0, 0],\n                  [0, cos(pitch), -sin(pitch)],\n                  [0, sin(pitch), cos(pitch)]])\n    R = np.array([[cos(roll), -sin(roll), 0],\n                  [sin(roll), cos(roll), 0],\n                  [0, 0, 1]])\n    return np.dot(Y, np.dot(P, R))\n#############\ndef draw_line(image, points):\n    color = (255, 0, 0)\n    cv2.line(image, tuple(points[0][:2]), tuple(points[3][:2]), color, 16)\n    cv2.line(image, tuple(points[0][:2]), tuple(points[1][:2]), color, 16)\n    cv2.line(image, tuple(points[1][:2]), tuple(points[2][:2]), color, 16)\n    cv2.line(image, tuple(points[2][:2]), tuple(points[3][:2]), color, 16)\n    return image\n\n\ndef draw_points(image, points):\n    for (p_x, p_y, p_z) in points:\n        cv2.circle(image, (p_x, p_y), int(1000 / p_z), (0, 255, 0), -1)\n    return image\n#############\ndef visualize(img, coords):\n    # You will also need functions from the previous cells\n    x_l = 1.02\n    y_l = 0.80\n    z_l = 2.31\n    \n    img = img.copy()\n    for point in coords:\n        # Get values\n        x, y, z = point['x'], point['y'], point['z']\n        yaw, pitch, roll = -point['pitch'], -point['yaw'], -point['roll']\n        # Math\n        Rt = np.eye(4)\n        t = np.array([x, y, z])\n        Rt[:3, 3] = t\n        Rt[:3, :3] = euler_to_Rot(yaw, pitch, roll).T\n        Rt = Rt[:3, :]\n        P = np.array([[x_l, -y_l, -z_l, 1],\n                      [x_l, -y_l, z_l, 1],\n                      [-x_l, -y_l, z_l, 1],\n                      [-x_l, -y_l, -z_l, 1],\n                      [0, 0, 0, 1]]).T\n        img_cor_points = np.dot(camera_matrix, np.dot(Rt, P))\n        img_cor_points = img_cor_points.T\n        img_cor_points[:, 0] /= img_cor_points[:, 2]\n        img_cor_points[:, 1] /= img_cor_points[:, 2]\n        img_cor_points = img_cor_points.astype(int)\n        # Drawing\n        img = draw_line(img, img_cor_points)\n        img = draw_points(img, img_cor_points[-1:])\n    \n    return img\n\n\ndef _regr_preprocess(regr_dict, flip=False):\n    if flip:\n        for k in ['x', 'pitch', 'roll']:\n            regr_dict[k] = -regr_dict[k]\n    for name in ['x', 'y', 'z']:\n        regr_dict[name] = regr_dict[name] / 100\n    regr_dict['roll'] = rotate(regr_dict['roll'], np.pi)\n    regr_dict['pitch_sin'] = sin(regr_dict['pitch'])\n    regr_dict['pitch_cos'] = cos(regr_dict['pitch'])\n    regr_dict.pop('pitch')\n    regr_dict.pop('id')\n    return regr_dict\n\ndef _regr_back(regr_dict):\n    for name in ['x', 'y', 'z']:\n        regr_dict[name] = regr_dict[name] * 100\n    regr_dict['roll'] = rotate(regr_dict['roll'], -np.pi)\n    \n    pitch_sin = regr_dict['pitch_sin'] / np.sqrt(regr_dict['pitch_sin']**2 + regr_dict['pitch_cos']**2)\n    pitch_cos = regr_dict['pitch_cos'] / np.sqrt(regr_dict['pitch_sin']**2 + regr_dict['pitch_cos']**2)\n    regr_dict['pitch'] = np.arccos(pitch_cos) * np.sign(pitch_sin)\n    return regr_dict\n\ndef preprocess_image(img, flip=False):\n    img = img[img.shape[0] // 2:]\n    bg = np.ones_like(img) * img.mean(1, keepdims=True).astype(img.dtype)\n    bg = bg[:, :img.shape[1] // 6]\n    img = np.concatenate([bg, img, bg], 1)\n    img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))\n    if flip:\n        img = img[:,::-1]\n    return (img / 255).astype('float32')\n\ndef get_mask_and_regr(img, labels, flip=False):\n    mask = np.zeros([IMG_HEIGHT // MODEL_SCALE, IMG_WIDTH // MODEL_SCALE], dtype='float32')\n    regr_names = ['x', 'y', 'z', 'yaw', 'pitch', 'roll']\n    regr = np.zeros([IMG_HEIGHT // MODEL_SCALE, IMG_WIDTH // MODEL_SCALE, 7], dtype='float32')\n    coords = str2coords(labels)\n    xs, ys = get_img_coords(labels)\n    for x, y, regr_dict in zip(xs, ys, coords):\n        x, y = y, x\n        x = (x - img.shape[0] // 2) * IMG_HEIGHT / (img.shape[0] // 2) / MODEL_SCALE\n        x = np.round(x).astype('int')\n        y = (y + img.shape[1] // 6) * IMG_WIDTH / (img.shape[1] * 4/3) / MODEL_SCALE\n        y = np.round(y).astype('int')\n        if x >= 0 and x < IMG_HEIGHT // MODEL_SCALE and y >= 0 and y < IMG_WIDTH // MODEL_SCALE:\n            mask[x, y] = 1\n            regr_dict = _regr_preprocess(regr_dict, flip)\n            regr[x, y] = [regr_dict[n] for n in sorted(regr_dict)]\n    if flip:\n        mask = np.array(mask[:,::-1])\n        regr = np.array(regr[:,::-1])\n    return mask, regr\n\ndef imread(path, fast_mode=False):\n    img = cv2.imread(path)\n    if not fast_mode and img is not None and len(img.shape) == 3:\n        img = np.array(img[:, :, ::-1])\n    return img\n\nimg = imread(PATH + 'train_images/ID_8a6e65317' + '.jpg')\nIMG_SHAPE = img.shape\n\ndef convert_3d_to_2d(x, y, z, fx = 2304.5479, fy = 2305.8757, cx = 1686.2379, cy = 1354.9849):\n    # stolen from https://www.kaggle.com/theshockwaverider/eda-visualization-baseline\n    return x * fx / z + cx, y * fy / z + cy\n\ndef optimize_xy(r, c, x0, y0, z0):\n    def distance_fn(xyz):\n        x, y, z = xyz\n        x, y = convert_3d_to_2d(x, y, z0)\n        y, x = x, y\n        x = (x - IMG_SHAPE[0] // 2) * IMG_HEIGHT / (IMG_SHAPE[0] // 2) / MODEL_SCALE\n        x = np.round(x).astype('int')\n        y = (y + IMG_SHAPE[1] // 6) * IMG_WIDTH / (IMG_SHAPE[1] *4/3) / MODEL_SCALE\n        y = np.round(y).astype('int')\n        return (x-r)**2 + (y-c)**2\n    \n    res = minimize(distance_fn, [x0, y0, z0], method='Powell')\n    x_new, y_new, z_new = res.x\n    return x_new, y_new, z0\n\ndef clear_duplicates(coords):\n    for c1 in coords:\n        xyz1 = np.array([c1['x'], c1['y'], c1['z']])\n        for c2 in coords:\n            xyz2 = np.array([c2['x'], c2['y'], c2['z']])\n            distance = np.sqrt(((xyz1 - xyz2)**2).sum())\n            if distance < DISTANCE_THRESH_CLEAR:\n                if c1['confidence'] < c2['confidence']:\n                    c1['confidence'] = -1\n    return [c for c in coords if c['confidence'] > 0]\n\ndef extract_coords(prediction):\n    logits = prediction[0]\n    regr_output = prediction[1:]\n    points = np.argwhere(logits > 0)\n    col_names = sorted(['x', 'y', 'z', 'yaw', 'pitch_sin', 'pitch_cos', 'roll'])\n    coords = []\n    for r, c in points:\n        regr_dict = dict(zip(col_names, regr_output[:, r, c]))\n        coords.append(_regr_back(regr_dict))\n        coords[-1]['confidence'] = 1 / (1 + np.exp(-logits[r, c]))\n        coords[-1]['x'], coords[-1]['y'], coords[-1]['z'] = optimize_xy(r, c, coords[-1]['x'], coords[-1]['y'], coords[-1]['z'])\n    coords = clear_duplicates(coords)\n    return coords\n\ndef coords2str(coords, names=['yaw', 'pitch', 'roll', 'x', 'y', 'z', 'confidence']):\n    s = []\n    for c in coords:\n        for n in names:\n            s.append(str(c.get(n, 0)))\n    return ' '.join(s)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Instead of 2 tensors *mask(1, 60, 192)* and *regr(1, 7, 60, 192)*. Modify dataset object to use single mask output with tensor of size *(1, 8, 60, 192)*\n\nThis is done by merging mask with regr using\n>mask = np.concatenate((np.expand_dims(mask, axis=0), regr), axis=0)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CarDataset(Dataset):\n    \"\"\"Car dataset.\"\"\"\n\n    def __init__(self, dataframe, root_dir, training=True, transform=None):\n        self.df = dataframe\n        self.root_dir = root_dir\n        self.transform = transform\n        self.training = training\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        \n        # Get image name\n        idx, labels = self.df.values[idx]\n        img_name = self.root_dir.format(idx)\n        \n        # Augmentation\n        flip = False\n        if self.training:\n            flip = np.random.randint(2) == 1\n        \n        # Read image\n        img0 = imread(img_name, True)\n        img = preprocess_image(img0, flip=flip)\n        img = np.rollaxis(img, 2, 0)\n        \n        # Get mask and regression maps\n        mask, regr = get_mask_and_regr(img0, labels, flip=False)\n        regr = np.rollaxis(regr, 2, 0)\n        \n        #Modify Mask to incorprorate regr \n        mask = np.concatenate((np.expand_dims(mask, axis=0), regr), axis=0)\n        \n        return [img, mask]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Test/Train Image directories\ntrain_images_dir = PATH + 'train_images/{}.jpg'\ntest_images_dir = PATH + 'test_images/{}.jpg'\n\n#Dataframe for train, valid and test\ndf_train, df_dev = train_test_split(train, test_size=0.01, random_state=42)\ndf_test = test","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Create dataset objects\ntrain_dataset = CarDataset(df_train, train_images_dir, training=True)\ndev_dataset = CarDataset(df_dev, train_images_dir, training=False)\n\nStop_tr = len(train_dataset)\nRandom_tr_ListOfIntegers = [random.randrange(0, Stop_tr) for iter in range(tr_limit)]\n\n#Limit Train dataset to 400 images\nif not full:\n    train_dataset = torch.utils.data.Subset(train_dataset, Random_tr_ListOfIntegers)\n\n# Create data generators - they will produce batches\ntrain_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE//2, shuffle=True, num_workers=0)\ndev_loader = DataLoader(dataset=dev_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n","execution_count":null,"outputs":[]},{"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\n        #  would be a nice idea if the upsampling could be learned too,\n        #  but my machine do not have enough memory to handle all those weights\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\n        self.conv = double_conv(in_ch, out_ch)\n\n    def forward(self, x1, x2=None):\n        x1 = self.up(x1)\n        \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        \n        # for padding issues, see \n        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n        \n        if x2 is not None:\n            x = torch.cat([x2, x1], dim=1)\n        else:\n            x = x1\n        x = self.conv(x)\n        return x\n\ndef get_mesh(batch_size, shape_x, shape_y):\n    mg_x, mg_y = np.meshgrid(np.linspace(0, 1, shape_y), np.linspace(0, 1, shape_x))\n    mg_x = np.tile(mg_x[None, None, :, :], [batch_size, 1, 1, 1]).astype('float32')\n    mg_y = np.tile(mg_y[None, None, :, :], [batch_size, 1, 1, 1]).astype('float32')\n    mesh = torch.cat([torch.tensor(mg_x).to(device), torch.tensor(mg_y).to(device)], 1)\n    return mesh\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model=xresnet50()\nmodel = nn.Sequential(*list(model.children())[:-4])\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CentResnet1(nn.Module):\n    '''Mixture of previous classes'''\n    def __init__(self, n_classes):\n        super(CentResnet1, self).__init__()\n        self.base_model = model\n        \n        # Lateral layers convert resnet outputs to a common feature size\n        self.lat8 = nn.Conv2d(512, 256, 1)\n        self.lat16 = nn.Conv2d(1024, 256, 1)\n        self.lat32 = nn.Conv2d(2048, 256, 1)\n        self.bn8 = nn.GroupNorm(16, 256)\n        self.bn16 = nn.GroupNorm(16, 256)\n        self.bn32 = nn.GroupNorm(16, 256)\n\n       \n        self.conv0 = double_conv(5, 64)\n        self.conv1 = double_conv(64, 128)\n        self.conv2 = double_conv(128, 512)\n        self.conv3 = double_conv(512, 1024)\n        \n        self.mp = nn.MaxPool2d(2)\n        \n        self.up1 = up(1282 , 512) #+ 1024\n        self.up2 = up(512 + 512, 256)\n        self.outc = nn.Conv2d(256, n_classes, 1)\n        \n    \n    def forward(self, x):\n        batch_size = x.shape[0]\n        mesh1 = get_mesh(batch_size, x.shape[2], x.shape[3])\n        x0 = torch.cat([x, mesh1], 1)\n        x1 = self.mp(self.conv0(x0))\n        x2 = self.mp(self.conv1(x1))\n        x3 = self.mp(self.conv2(x2))\n        x4 = self.mp(self.conv3(x3))\n        \n                # Run frontend network\n        feats32 = self.base_model(x)\n        lat32 = F.relu(self.bn32(self.lat32(feats32)))\n        \n        # Add positional info\n        mesh2 = get_mesh(batch_size, lat32.shape[2], lat32.shape[3])\n        feats = torch.cat([lat32, mesh2], 1)\n        x = self.up1(feats, x4)\n        x = self.up2(x, x3)\n        x = self.outc(x)\n        return x\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The prediction of size (1, 8, 60, 192) has to be split to regr and mask parts to facilitate loss calculation. We change back the mask to (1, 60, 192) and regr to (1, 7, 60, 192) using this:\n>regr = mask[0,1:].unsqueeze(0)\n\n>mask = mask[0,0:1]"},{"metadata":{"trusted":true},"cell_type":"code","source":"def mod_criterion(prediction, mask, weight=0.4, size_average=True):\n    #Extract regr and mask values from the combined mask\n    regr = mask[0,1:].unsqueeze(0)\n    mask = mask[0,0:1]\n    # Binary mask loss\n    pred_mask = torch.sigmoid(prediction[:, 0])\n#     mask_loss = mask * (1 - pred_mask)**2 * torch.log(pred_mask + 1e-12) + (1 - mask) * pred_mask**2 * torch.log(1 - pred_mask + 1e-12)\n    mask_loss = mask * torch.log(pred_mask + 1e-12) + (1 - mask) * torch.log(1 - pred_mask + 1e-12)\n    mask_loss = -mask_loss.mean(0).sum()\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 = weight*mask_loss +(1-weight)* regr_loss\n    if not size_average:\n        loss *= prediction.shape[0]\n    return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Import FastaiV2 references, create learner and proceed as usual."},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision import datasets, transforms\nfrom torch.optim.lr_scheduler import StepLR\nfrom fastai2.optimizer import Adam\n\ndata = DataBunch(train_loader, dev_loader).cuda()\nlearn = Learner(data, CentResnet1(8), loss_func=mod_criterion, opt_func=Adam)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.lr_find()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.fit_one_cycle(n_epochs, 1e-2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.recorder.plot_loss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.save(\"Centernet-Xresnet50-ph1\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Please upvote if you found the notebook to be useful."},{"metadata":{},"cell_type":"markdown","source":"References:\n* https://nbviewer.jupyter.org/github/fastai/fastai2/blob/master/nbs/migrating.ipynb\n* https://www.kaggle.com/hocop1/centernet-baseline\n* https://www.kaggle.com/phoenix9032/center-resnet-starter"}],"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":1}