{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"\"\"\"\nI'll only be covering the detection of ships in this script, there won't be any segmentation stuff.\n\nThis will also not be an information session. The purpose is not to teach but to introduce a new tool\nthe learning will have to be done on the individual level.\n\"\"\"\nimport os\nimport tqdm\n\nimport numpy as np\nimport pandas as pd\n\nfrom PIL import Image, ImageDraw\n\nimport torch\nimport torchvision\n\nINPUT_FOLDER = '../input'\nTRAIN_FOLDER = os.path.join(INPUT_FOLDER, 'train')\nTEST_FOLDER  = os.path.join(INPUT_FOLDER, 'test')\nTRAIN_SEGMENTATION_FILE = os.path.join(INPUT_FOLDER, 'train_ship_segmentations.csv')\n\nIMAGE_SHAPE = [768, 768]\n\nBAD_IMAGES = ['6384c3e78.jpg']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"4b392447e6822c51d72d7e4d200425459db37803"},"cell_type":"code","source":"def rle_decode(mask_rle, shape):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return\n    Returns numpy array, 1 - mask, 0 - background\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape).T\n\ndef mask_to_bbox(mask):\n    img_h, img_w = mask.shape[:2]\n\n    rows = np.any(mask, axis=1)\n    cols = np.any(mask, axis=0)\n\n    rmin, rmax = np.where(rows)[0][[0, -1]]\n    cmin, cmax = np.where(cols)[0][[0, -1]]\n\n    x1 = int(max(cmin - 1, 0))\n    y1 = int(max(rmin - 1, 0))\n    x2 = int(min(cmax + 1, img_w))\n    y2 = int(min(rmax + 1, img_h))\n\n    return x1, y1, x2, y2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"050553eec9c841d7f67674a6e5101bd95e73e32a"},"cell_type":"code","source":"# First to perform some preprocessing, I suggest skipping this code block and move on to the next see what I'm after\ntrain_df = pd.read_csv(TRAIN_SEGMENTATION_FILE)\ntrain_df_values = train_df.values\ntrain_image_to_ann = {}\nfor entry in tqdm.tqdm(train_df_values, 'Create image to ann table'):\n    image_file = entry[0]\n    segmentation = entry[1]\n    \n    if image_file in BAD_IMAGES or pd.isnull(segmentation):\n        continue\n        \n    # Calc bbox\n    mask = rle_decode(segmentation, IMAGE_SHAPE)\n    bbox = mask_to_bbox(mask)\n    ann = np.array(bbox + (0,))\n    \n    if image_file in train_image_to_ann:\n        train_image_to_ann[image_file].extend([ann])\n    else:\n        train_image_to_ann[image_file] = [ann]\n\nfor image_file, anns in train_image_to_ann.items():\n    train_image_to_ann[image_file] = np.array(anns, dtype='float32')\n\nimage_files = list(train_image_to_ann.keys())\n\n# Clean up for freeing of memory in the future\ndel train_df\ndel train_df_values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"83df5633d7b8e14d1ee988e6d62570c5d19cc38c"},"cell_type":"code","source":"# The data format I wanted to make is an image to bounding boxes dataset\n# The labels for each image will look a little something like this\n# Each row represents a bounding box around a ship\n# and the values represent x1, y1, x2, y2, class_id\n# If you don't know what those means, go google and unstuck yourself\ntrain_image_to_ann[image_files[5]]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"27e56a5fe5fa1ce16d2edcf89201c38b17c73a04"},"cell_type":"code","source":"# To further visualize what I'm referring to, let's plot an image with the labels out\nsample_img = Image.open(os.path.join(TRAIN_FOLDER, image_files[5]))\ndraw = ImageDraw.Draw(sample_img)\nfor ann in train_image_to_ann[image_files[5]]:\n    draw.rectangle(ann[:4], outline=(255, 0, 0))\nsample_img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"fabf71b23ddcfc7e3b58eea304d77d44c3bd3666"},"cell_type":"code","source":"# With this dataset, I'm going to make a dataloader\nclass DetectionDataset(torch.utils.data.Dataset):\n    def __init__(self, image_to_ann):\n        self.image_to_ann = image_to_ann\n        self.image_files = list(image_to_ann.keys())\n        self.image_to_tensor = torchvision.transforms.Compose([\n            torchvision.transforms.ToTensor(),\n            torchvision.transforms.Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225]\n            )\n        ])\n    \n    def __len__(self):\n        return len(self.image_to_ann)\n\n    def __getitem__(self, idx):\n        image_file = self.image_files[idx]\n        image = Image.open(os.path.join(TRAIN_FOLDER, image_file))\n        image = self.image_to_tensor(image)\n        anns = self.image_to_ann[image_file]\n        anns = torch.from_numpy(anns)\n        return image, anns\n\ndataset = DetectionDataset(train_image_to_ann)\ndataset_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=2, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"d35067cf62777ba0ab98b3791992db38eae79b14"},"cell_type":"code","source":"for batch in dataset_loader:\n    break\n\n# We now see that batches of data produced by datadetloader produces batches of image, annotations pair\nbatch[0].shape, batch[1].shape","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"# Now to introduce something new, we first install a detection package\n!rm -rf torch_collections/\n!wget -q 'https://github.com/mingruimingrui/torch-collections/archive/0.4b.zip' -O torch-collections.zip\n!unzip -oq torch-collections.zip\n!mv torch-collections-0.4b/torch_collections .\n!rm -rf torch-collections-0.4b/  torch-collections.zip","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"e83d84cfcd887024651fef283ae4dd985fa0361b"},"cell_type":"code","source":"# And follow the package use case\nfrom torch_collections import RetinaNet\n\n# Build a model\nmodel = RetinaNet(1).train().cuda()\n\n# And an optimizer\noptimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.00001)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"92d38c1c7377b3630e751dfafad0e3612dd3cb6c"},"cell_type":"code","source":"# And we train the model\nmax_count = 6000\ncount = 0\npbar = tqdm.tqdm(total=max_count, desc='training model')\nfor batch in dataset_loader:\n    pbar.update(1)\n    optimizer.zero_grad()\n    \n    count\n    \n    batch_image = batch[0].cuda()\n    batch_annotations = batch[1].cuda()\n    \n    loss = model(batch_image, batch_annotations)\n    \n    if loss is not None:\n        # loss can be none when no valid anchors are available\n        loss.backward()\n        optimizer.step()\n    \n    del batch_image\n    del batch_annotations\n    del batch\n    \n    count += 1\n    if count >= max_count:\n        break\n    \npbar.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"009a7d75177d124b3099ce7ff3369eaa0298d481"},"cell_type":"code","source":"# Now let's look at the results shall we?\nmodel = model.eval()\n\nviz = []\nmax_viz = 5\ncount_viz = 0\n\nfor image_file in os.listdir(TEST_FOLDER):\n    image = Image.open(os.path.join(TEST_FOLDER, image_file))\n    image_tensor = dataset.image_to_tensor(image)\n    image_tensor = image_tensor.unsqueeze(0).cuda()\n\n    dets = model(image_tensor)[0]\n    \n    scores = dets['scores'].cpu().data.numpy()\n    boxes = dets['boxes'].cpu().data.numpy()[scores > 0.7]\n    \n    if len(boxes) > 0:\n        boxes = boxes.round()\n        \n        image = Image.open(os.path.join(TEST_FOLDER, image_file))\n        draw = ImageDraw.Draw(image)\n        for box in boxes:\n            draw.rectangle(box[:4], outline=(255, 0, 0))\n        \n        count_viz += 1\n        viz.extend([image])\n    \n    del image_tensor\n    del dets\n        \n    if count_viz >= max_viz:\n        break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"09172f83f0bf199daadcafd0a9c96be3106ef3b8"},"cell_type":"code","source":"viz[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"b176831d0d100346bc4968f54b79b905d8e5e771"},"cell_type":"code","source":"viz[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"49b92c303bbfd7ce91ec059fa3c7dba9e5068e48"},"cell_type":"code","source":"viz[2]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"ebe4db300f8148089a18c64319efdbb91a669bd3"},"cell_type":"code","source":"viz[3]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_uuid":"61ac21eff2ca384ebbc21cf7049d88628283d652"},"cell_type":"code","source":"viz[4]","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}