{"cells":[{"metadata":{},"cell_type":"markdown","source":"# OpenImage Challenge 2019\n\n## training fasterrcnn in pytorch\n\nfiles (N is [0-9A-F]):\n\nBASE_DIR/classes-segmentation.txt   …download from https://storage.googleapis.com/openimages/v5/classes-segmentation.txt\n\nBASE_DIR/train-images-N/*.jpg   …training image from s3://open-images-dataset/tar/train_N.tar.gz\n\nBASE_DIR/mask-images-N/*.png   …mask image from https://storage.googleapis.com/openimages/v5/train-masks/train-masks-N.zip\n\nTEST_DIR/*.jpg   …test image for prediction\n\ntemporary directory:\n\nTEMP_DIR/join-masks/\n\nTEMP_DIR/output-images/"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom tqdm import tqdm_notebook as tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"TEMP_DIR = './tmp/'\nBASE_DIR = '../input/openimage2019segmentationsubset2000/subset2000/'\nTEST_DIR = '../input/open-images-2019-instance-segmentation/test/'\n\nwith open(BASE_DIR+\"classes-segmentation.txt\") as f:\n    CLASSES = [c.strip() for c in f.readlines()]\nCLASSES = [\"__background__\"] + CLASSES\nNUM_CLASS = len(CLASSES)\nCLOP_SIZE = 480\nBASE_SIZE = 520\nNUM_CROP = 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"BATCH_SIZE = 2\nNUM_WORKERS = 3\nNUM_EPOCHS = 2\nNUM_GPUS = 1\n# Use all data: 0-f, z is subset2k\nUSING_SPLITS = [\"z\"] #[\"0\",\"1\",\"2\",\"3\",\"4\",\"5\",\"6\",\"7\",\"8\",\"9\",\"a\",\"b\",\"c\",\"d\",\"e\",\"f\"]","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if not os.path.isdir(TEMP_DIR):\n    os.mkdir(TEMP_DIR)\n    os.mkdir(TEMP_DIR+\"join-masks\")\n    os.mkdir(TEMP_DIR+\"output-images\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def _mask_filepart_classname(name):\n    if name.startswith(\"m\"):\n        return \"/m/\" + name[1:]\n    return name","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"import math\ndef _create_resize_image(img, ismask=False, tosize=None):\n    long_side_size = BASE_SIZE * 2\n    if img.height < img.width:\n        scale = img.width / long_side_size\n        size = (long_side_size, max(BASE_SIZE,math.ceil(img.height / scale)))\n    else:\n        scale = img.height / long_side_size\n        size = (max(BASE_SIZE,math.ceil(img.width / scale)), long_side_size)\n    return img.resize(size if tosize is None else tosize, Image.NEAREST if ismask else Image.BILINEAR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"import pickle\ndef _make_openimage2019_mask(split_name):\n    img_paths = []\n    mask_paths = []\n    img_folder = os.path.join(BASE_DIR, 'train-images-'+split_name)\n    mask_folder = os.path.join(BASE_DIR, 'mask-images-'+split_name)\n    join_folder = os.path.join(TEMP_DIR, 'join-masks')\n    img_folder_list = sorted(list(os.listdir(img_folder)))\n    image_mask = {}\n    for filename in os.listdir(mask_folder):\n        basename, _ = os.path.splitext(filename)\n        maskname = basename.split(\"_\")\n        if filename.endswith(\".png\"):\n            imgpath = os.path.join(img_folder, filename)\n            imagename = maskname[0] + '.jpg'\n            imagepath = os.path.join(img_folder, imagename)\n            if os.path.isfile(imagepath):\n                if imagepath not in image_mask:\n                    image_mask[imagename] = [filename]\n                else:\n                    image_mask[imagename].append(filename)\n            else:\n                print('cannot find the image:', imagepath)\n\n    for imagename, masknames in tqdm(image_mask.items()):\n        for nc in range(NUM_CROP):\n            imgpath = os.path.join(img_folder, imagename)\n            basename, _ = os.path.splitext(imagename)\n            joinpath = os.path.join(join_folder, basename+\"-\"+str(nc)+\".pkl\")\n            if os.path.isfile(joinpath):\n                continue\n\n            img_rs = _create_resize_image(Image.open(imgpath)).convert('RGB')\n            \n            crop_x = np.random.randint(img_rs.width-CLOP_SIZE)\n            crop_y = np.random.randint(img_rs.height-CLOP_SIZE)\n            img = img_rs.crop((crop_x, crop_y, crop_x+CLOP_SIZE, crop_y+CLOP_SIZE))\n\n            boxes = []\n            masks = []\n            labels = []\n\n            for filename in masknames:\n                basename, _ = os.path.splitext(filename)\n                maskname = basename.split(\"_\")\n                maskpath = os.path.join(mask_folder, filename)\n                maskflag = _create_resize_image(Image.open(maskpath), ismask=True, tosize=(img_rs.width,img_rs.height))\n                maskflag = maskflag.crop((crop_x, crop_y, crop_x+CLOP_SIZE, crop_y+CLOP_SIZE))\n                maskflag = np.array(maskflag.convert('1'))\n                maskclass = _mask_filepart_classname(maskname[1])\n                if np.sum(maskflag) > 0 and maskclass in CLASSES:\n                    labels.append(CLASSES.index(maskclass))\n                    pos = np.where(maskflag)\n                    xmin = np.min(pos[1])\n                    xmax = np.max(pos[1])\n                    ymin = np.min(pos[0])\n                    ymax = np.max(pos[0])\n                    boxes.append([xmin, ymin, xmax, ymax])\n                    masks.append(maskflag)\n\n            if len(boxes) > 0:\n\n                boxes = np.array(boxes)\n                masks = np.array(masks)\n                labels = np.array(labels)\n\n                idx = 0\n                if imagename in img_folder_list:\n                    idx = img_folder_list.index(imagename)\n                image_id = [idx]\n\n                if boxes.shape[0] == 0:\n                    area = 0\n                else:\n                    area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n\n                iscrowd = np.zeros((len(boxes),), dtype=np.int64)\n\n                target = {}\n                target[\"boxes\"] = boxes\n                target[\"labels\"] = labels\n                target[\"masks\"] = masks\n                target[\"image_id\"] = image_id\n                target[\"area\"] = area\n                target[\"iscrowd\"] = iscrowd\n\n                imgf = np.array(img, dtype=np.float32).transpose(2,0,1)\n\n                with open(joinpath, 'wb') as f:\n                    pickle.dump((imgf,target), f)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"for z in USING_SPLITS:\n    _make_openimage2019_mask(z)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"join_folder = os.path.join(TEMP_DIR, 'join-masks')\njoin_files = sorted([f for f in os.listdir(join_folder) if f.endswith(\".pkl\")])","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"class MyDataset(object):\n    def __init__(self):\n        self.join_folder = os.path.join(TEMP_DIR, 'join-masks')\n        self.files = sorted([f for f in os.listdir(self.join_folder) if f.endswith(\".pkl\")])\n\n    def __getitem__(self, idx):\n        mask_path = os.path.join(self.join_folder, self.files[idx])\n        with open(mask_path, 'rb') as f:\n            imgf,target = pickle.load(f)\n        target[\"boxes\"] = torch.as_tensor(target[\"boxes\"], dtype=torch.float32)\n        target[\"labels\"] = torch.as_tensor(target[\"labels\"], dtype=torch.int64)\n        target[\"masks\"] = torch.as_tensor(target[\"masks\"], dtype=torch.uint8)\n        target[\"image_id\"] = torch.as_tensor(target[\"image_id\"], dtype=torch.int32)\n        target[\"area\"] = torch.as_tensor(target[\"area\"], dtype=torch.float32)\n        target[\"iscrowd\"] = torch.as_tensor(target[\"iscrowd\"], dtype=torch.int64)\n\n        imgf = torch.as_tensor(imgf, dtype=torch.float32)\n \n        return imgf, target[\"boxes\"], target[\"labels\"], target[\"masks\"], target[\"image_id\"], target[\"area\"], target[\"iscrowd\"]\n\n    def __len__(self):\n        return len(self.files)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"import torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n\nrcnnmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, num_classes=NUM_CLASS, pretrained_backbone=True)\nrcnnmodel.eval()\nmodel = torch.nn.DataParallel(rcnnmodel)\nmodel.cuda()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"import torchvision.transforms as T\n\ndef get_transform(train):\n    transforms = []\n    transforms.append(T.ToTensor())\n    if train:\n        transforms.append(T.RandomHorizontalFlip(0.5))\n    return T.Compose(transforms)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"import math\nimport sys\nimport time\nimport torch\n\ndef train_one_epoch(model, optimizer, data_loader, device, epoch):\n    model.train()\n\n    lr_scheduler = None\n    if epoch == 0:\n        warmup_factor = 1. / 1000\n        warmup_iters = min(1000, len(data_loader) - 1)\n\n    prog = tqdm(data_loader, total=len(data_loader))\n    for imgfs, boxes, labels, masks, image_id, area, iscrowd in prog:\n        images = []\n        targets = []\n\n        for i in range(imgfs.shape[0]):\n            images.append(imgfs[i].cuda())\n            target = {}\n            target[\"boxes\"] = boxes[i].cuda()\n            target[\"labels\"] = labels[i].cuda()\n            target[\"masks\"] = masks[i].cuda()\n            target[\"image_id\"] = image_id[i].cuda()\n            target[\"area\"] = area[i].cuda()\n            target[\"iscrowd\"] = iscrowd[i].cuda()\n            targets.append(target)\n\n        loss_dict = model(images, targets)\n        losses = sum(loss for loss in loss_dict.values())\n\n        prog.set_description(\"loss:%03f\"%losses)\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"device = torch.device('cuda')\n\ndataset = MyDataset()\n\ndata_loader = torch.utils.data.DataLoader(\n    dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS)\n\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.Adam(params)\n\nfor epoch in range(NUM_EPOCHS):\n    train_one_epoch(model, optimizer, data_loader, device, epoch)\n    torch.save(rcnnmodel.state_dict(), \"checkpoint-%d\"%epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"torch.save(rcnnmodel.state_dict(), \"final_model\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.rmtree(TEMP_DIR)","execution_count":null,"outputs":[]}],"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.9"}},"nbformat":4,"nbformat_minor":1}