{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Purpose\nIn order to efficiently train a network, a good data generator is needed. One that does not exceed ram limit, is efficient and fast.\nHere I present several data generating classes along with some helper methods to achieve that. \n\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"TRAIN_IMG_IDS = ['095bf7a1f', 'afa5e8098', 'e79de561c', '4ef6695ce', '0486052bb', 'b9a3865fc', 'aaa6a05cc', '2f6ecfcdf',\n                 '54f2eec69', 'c68fe75ea', '8242609fa', '1e2425f28', 'b2dc8411c', 'cb2d976f4', '26dc41664']\n\nTEST_IMG_IDS = ['aa05346ff', '2ec3f1bb9', '3589adb90', 'd488c759a', '57512b7f1']\n\nDATA_DIR = \"/kaggle/input/hubmap-kidney-segmentation\"\n\n\nimport json\nimport os\nimport random\n\nimport tifffile as tiff\nimport cv2\nfrom PIL import Image\nfrom keras_preprocessing.image import load_img\nfrom matplotlib.patches import Rectangle\nfrom tensorflow import keras\nimport pandas as pd\n\nimport numpy as np\nfrom matplotlib import pyplot as plt\n\n#%%\n\n\ndef scale_down(image, scale):\n    nx = (image.shape[0] // scale) * scale\n    ny = (image.shape[1] // scale) * scale\n    return cv2.resize(image[:nx, :ny], (ny // scale, nx // scale))\n\n\ndef check_if_in_dataset(filename):\n    if filename in TRAIN_IMG_IDS:\n        return \"train\"\n    elif filename in TEST_IMG_IDS:\n        return \"test\"\n    else:\n        return \"\"\n\n\ndef show_hsv(im):\n    hsv_img = cv2.cvtColor(im, cv2.COLOR_RGB2HSV)\n    l = [\"H\", \"S\", \"V\"]\n    plt.figure()\n    for i in range(3):\n        count, bins = np.histogram(hsv_img[:,:,i], bins=np.arange(hsv_img.max()))\n        plt.plot( bins[1:-1], count[1:], label=l[i])\n    plt.legend()\n    plt.show()\n\n\ndef get_image(file_name, dataset, initial_scale, PATH=DATA_DIR):\n    image = tiff.imread('{}/{}/{}.tiff'.format(PATH,dataset, file_name))\n    dim = image.ndim\n\n    if dim == 5:\n        image = image[0, 0, :, :, :].transpose((1, 2, 0))\n    elif image.shape[0] == 3:\n        image = image.transpose((1, 2, 0))\n\n    return scale_down(image, initial_scale)\n\n\ndef get_pink_mask(image, mask_scale, blur_kernel):\n    rgb_img = scale_down(image, mask_scale)\n    hsv_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2HSV)\n    light_pink = (120, 10, 20)\n    dark_pink = (150, 190, 255)\n    mask_pink = cv2.inRange(hsv_img, light_pink, dark_pink)\n    blured = cv2.blur(mask_pink,blur_kernel)\n    return (blured > 127).astype(np.uint8)\n\n\ndef get_glomerous_polygon_list(file_name):\n\n    glomerulus = open('{}/{}/{}.json'.format(DATA_DIR,\"train\", file_name))\n    json_obj = json.load(glomerulus)\n\n    polygon_list = []\n\n    for elem in json_obj:\n        narr = np.array(elem['geometry']['coordinates'][0])\n        polygon_list.append(narr)\n\n    return polygon_list\n\ndef get_glomerulus_mask(filename, initial_scale, PATH=DATA_DIR):\n    df = pd.read_csv(PATH + \"/HuBMAP-20-dataset_information.csv\")\n    w = df[df.image_file == (filename + \".tiff\")].width_pixels.tolist()[0]\n    h = df[df.image_file == (filename + \".tiff\")].height_pixels.tolist()[0]\n    image_shape = (h,w)\n    mask = np.zeros(image_shape)\n    polygon_list = get_glomerous_polygon_list(filename)\n    color = 1\n    for coordinates in polygon_list:\n        cv2.fillPoly(mask, pts=[coordinates.astype(np.int32)], color=color)\n    nx = (mask.shape[0] // initial_scale) * initial_scale\n    ny = (mask.shape[1] // initial_scale) * initial_scale\n    return cv2.resize(mask[:nx, :ny], (ny // initial_scale, nx // initial_scale))\n\n#%%\n\ndef get_mask_grid_sum(mask, mask_scale, grid_size, x, y):\n    x_from = x // mask_scale\n    y_from = y // mask_scale\n    x_to = (x + grid_size) // mask_scale\n    y_to = (y + grid_size) // mask_scale\n    return mask[x_from:x_to,y_from:y_to].sum()\n\n\ndef get_grids(image, mask, initial_scale, GRID_SIZE, OVERLAP, include_empty_grids=False):\n    mask_scale = 10\n    mask = scale_down(mask, mask_scale)\n    # mask_threshold determined by roughly single glomerulus size\n    mask_threshold = (400/initial_scale/mask_scale)**2\n    # blur range tested empiricaly to be about half of glomerulus size\n    blur_range = max(200 // initial_scale // mask_scale, 3)\n    img_shape = image.shape\n    pink_mask = get_pink_mask(image, mask_scale, (blur_range, blur_range))\n    grids = []\n    for row in range(0, img_shape[0] - GRID_SIZE, GRID_SIZE - OVERLAP):\n        for col in range(0, img_shape[1] - GRID_SIZE, GRID_SIZE - OVERLAP):\n            valid_pixels = get_mask_grid_sum(pink_mask,mask_scale,GRID_SIZE, row, col)\n            if valid_pixels > mask_threshold:\n                glomeruli_pixels = get_mask_grid_sum(mask,mask_scale,GRID_SIZE, row, col)\n                if glomeruli_pixels > 0 or include_empty_grids:\n                    grids.append([row,col,row+GRID_SIZE, col+GRID_SIZE])\n    return grids\n\n\ndef get_BB(coords):\n    ret = np.zeros((4,))\n    # ret: [min_x,min_y,max_x,max_y]\n    ret[:2] = np.min(coords, axis=0)[::-1]\n    ret[2:] = np.max(coords, axis=0)[::-1]\n    return ret\n\n\ndef in_grid(x, GRID_SQUARE_SIZE):\n    if x[0] >= GRID_SQUARE_SIZE or x[2] <= 0:\n        return False\n    if x[1] >= GRID_SQUARE_SIZE or x[3] <= 0:\n        return False\n    return True\n\n\ndef turnicate(bb, GRID_SQUARE_SIZE):\n    bb[bb < 0] = 0\n    bb[bb > (GRID_SQUARE_SIZE - 1)] = GRID_SQUARE_SIZE - 1\n    return bb\n\n\ndef get_BB_within(polys, x_start, y_start, GRID_SIZE, initial_scale):\n    GRID_SIZE = GRID_SIZE*initial_scale\n    x_start = x_start*initial_scale\n    y_start = y_start*initial_scale\n    offset = np.array([x_start, y_start, x_start, y_start])\n    ofs = [get_BB(coords) - offset for coords in polys]\n    relative = [turnicate(x, GRID_SIZE) for x in ofs if in_grid(x, GRID_SIZE)]\n    return [x//initial_scale for x in relative if in_grid(x, GRID_SIZE)]\n\n\n#%%\n\n#%%\n\ndef draw_img_boxes(img, boxes):\n    plt.figure()\n    plt.imshow(img)\n    ax = plt.gca()\n    mask = np.zeros(img.shape[:2])\n    for coords in boxes:\n        lx, ly, rx, ry = coords.astype(np.int)\n        mask[lx:rx,ly:ry] += 1\n    plt.imshow(mask, alpha=0.2)\n    for coords in boxes:\n        # flip coords for Rectangle\n        y,x = coords[0:2]\n        dy,dx = coords[2] - coords[0], coords[3] - coords[1]\n        rect = Rectangle((x,y),dx, dy,linewidth=2,edgecolor='w',facecolor='none')\n        ax.add_patch(rect)\n\n#%%\n\ndef showcase_bbs(file_name, grid_i):\n    im = get_image(\"afa5e8098\", \"train\", 2)\n    mask = get_glomerulus_mask(\"afa5e8098\", 2)\n    g = get_grids(im, mask, 2, 1024, 256)\n    lx, ly, rx, ry = g[grid_i]\n    polygon_list = get_glomerous_polygon_list(file_name)\n    bbs = get_BB_within(polygon_list, lx, ly, 1024, 2)\n    draw_img_boxes(im[lx:rx,ly:ry], bbs)\n#%%\n# uncomment to see img with bounding boxes\n# showcase_bbs(\"afa5e8098\", 43)\n\n#%%\n\ndef showcase_grid(file_name):\n    im = get_image(file_name, \"train\", 10)\n    mask = get_glomerulus_mask(file_name, 10)\n    g = get_grids(im,mask, 10, 256,64)\n\n    im = scale_down(im, 10)\n    mask = scale_down(mask, 10)\n\n    gm = np.zeros((im.shape[0]*10, im.shape[1]*10))\n    for lx,ly,rx,ry in g:\n        gm[lx:rx,ly:ry] += 1\n    gm = cv2.resize(gm, (gm.shape[1]//10,gm.shape[0]//10))\n    plt.figure()\n    plt.imshow(im)\n    plt.imshow(mask, alpha=0.3)\n    plt.imshow(gm, alpha=0.3)\n#%%\n# uncomment to see grid\n# showcase_grid(\"afa5e8098\")\n\n#%%\n\n\nclass MyDataSequenceSlow(keras.utils.Sequence):\n    \"\"\"Helper to iterate over the data (as Numpy arrays).\"\"\"\n\n    def __init__(self, batch_size, train_img_names, initial_scale=2, GRID_SIZE=1024, OVERLAP=256):\n        self.batch_size = batch_size\n        self.batches = []\n        self.initial_scale = initial_scale\n        print(\"starting preprocessing\")\n        for file_name in train_img_names:\n            print(\"Preparing\", file_name, \"...\")\n            image = get_image(file_name, \"train\", initial_scale)\n            mask = get_glomerulus_mask(file_name, initial_scale)\n            grids = get_grids(image, mask, file_name, initial_scale, GRID_SIZE, OVERLAP)\n            random.shuffle(grids)\n            print(\"Prepared\", len(grids), \"grids.\", len(grids)%batch_size, \"lost in batching.\")\n            for i in range(0, len(grids) - batch_size, batch_size):\n                self.batches.append((file_name, grids[i:i+batch_size]))\n            print(\"Batches saved.\")\n\n    def __len__(self):\n        return len(self.batches)\n\n    def __getitem__(self, idx):\n        \"\"\"Returns tuple (input, target) correspond to batch #idx.\"\"\"\n        raise Exception(\"Not implemented.\")\n\n\nclass MyDataSequenceSlowMask(MyDataSequenceSlow):\n\n    def __getitem__(self, idx):\n        file_name, grids = self.batches[idx]\n\n        image = get_image(file_name, \"train\", self.initial_scale)\n        mask = get_glomerulus_mask(file_name, self.initial_scale)\n\n        grid_imgs = []\n        grid_masks = []\n        for lx,ly,rx,ry in grids:\n            im = image[lx:rx,ly:ry]\n            grid_imgs.append(im)\n            m = mask[lx:rx,ly:ry]\n            grid_masks.append(m)\n        return np.array(grid_imgs), np.array(grid_masks)\n\n\nclass MyDataSequenceSaving(keras.utils.Sequence):\n\n    def __init__(self, batch_size, train_img_names, initial_scale=2, GRID_SIZE=1024, OVERLAP=256,preprocess=True):\n        self.CACHE_PATH = \"/kaggle/working/cache\"\n        self.TRAIN_PATH = self.CACHE_PATH + \"/train\"\n        self.IMAGES = self.TRAIN_PATH + \"/images\"\n        self.MASKS = self.TRAIN_PATH + \"/masks\"\n        self.GRID_SIZE = GRID_SIZE\n        self.batch_size = batch_size\n        self.img_list = []\n        self.initial_scale = initial_scale\n        if preprocess:\n            print(\"starting preprocessing\")\n            try:\n                os.makedirs(self.IMAGES)\n                os.makedirs(self.MASKS)\n            except:\n                print(\"dirs already exist\")\n\n            for file_name in train_img_names:\n                bb_dict = {file_name: []}\n                print(\"Preparing\", file_name, \"...\")\n                image = get_image(file_name, \"train\", initial_scale)\n                mask = get_glomerulus_mask(file_name, initial_scale)\n                grids = get_grids(image, mask, initial_scale, GRID_SIZE, OVERLAP)\n                polygon_list = get_glomerous_polygon_list(file_name)\n                print(\"Prepared\", len(grids), \"grids.\")\n                for i,(lx,ly,rx,ry) in enumerate(grids):\n                    im = Image.fromarray(image[lx:rx, ly:ry])\n                    im.save(f\"{self.IMAGES}/{file_name}_{i}.png\")\n\n                    msk = Image.fromarray(mask[lx:rx, ly:ry].astype(np.uint8))\n                    msk.save(f\"{self.MASKS}/{file_name}_{i}.png\")\n\n                    bbs = get_BB_within(polygon_list, lx,ly,GRID_SIZE, initial_scale)\n                    bbs_lists = [x.tolist() for x in bbs]\n                    bb_dict[file_name].append(bbs_lists)\n\n                    self.img_list.append(file_name + \"_\" + str(i))\n                with open(f'{self.TRAIN_PATH}/bbs-{file_name}.json', 'w') as fp:\n                    json.dump(bb_dict, fp)\n                print(\"pics saved.\")\n        else:\n            for dirname, _, filenames in os.walk(self.IMAGES):\n                for filename in filenames:\n                    self.img_list.append(filename[:-4])\n        random.shuffle(self.img_list)\n\n    def __len__(self):\n        return len(self.img_list)//self.batch_size\n\n    def __getitem__(self, idx):\n        \"\"\"Returns tuple (input, target) correspond to batch #idx.\"\"\"\n        raise Exception(\"Not implemented.\")\n\n\nclass MyDataSequenceSavingMask(MyDataSequenceSaving):\n\n    def __getitem__(self, idx):\n        \"\"\"Returns tuple (input, target) correspond to batch #idx.\"\"\"\n\n        batch = self.img_list[idx*self.batch_size:(idx+1)*self.batch_size]\n        grid_imgs = []\n        grid_masks = []\n\n        for img_name in batch:\n            img_path = f\"{self.IMAGES}/{img_name}.png\"\n            mask_path = f\"{self.MASKS}/{img_name}.png\"\n            image = load_img(img_path, target_size=(self.GRID_SIZE, self.GRID_SIZE))\n            mask = load_img(mask_path, target_size=(self.GRID_SIZE, self.GRID_SIZE), color_mode=\"grayscale\")\n\n            grid_imgs.append(np.array(image))\n            grid_masks.append(np.expand_dims(mask, axis=2))\n\n        return np.array(grid_imgs), np.array(grid_masks)\n\nclass MyDataSequenceSavingBBs(MyDataSequenceSaving):\n\n    def __getitem__(self, idx):\n        \"\"\"Returns tuple (input, target) correspond to batch #idx.\"\"\"\n\n        batch = self.img_list[idx*self.batch_size:(idx+1)*self.batch_size]\n        grid_imgs = []\n        grid_BBs = []\n\n        for img_name in batch:\n            file_name,img_idx = img_name.split(\"_\")\n            img_path = f\"{self.IMAGES}/{img_name}.png\"\n            bbs_path = f\"{self.TRAIN_PATH}/bbs-{file_name}.json\"\n            image = load_img(img_path, target_size=(self.GRID_SIZE, self.GRID_SIZE))\n            with open(bbs_path) as json_file:\n                bbs = json.load(json_file)\n            grid_imgs.append(np.array(image))\n            bbs_lists = bbs[file_name][int(img_idx)]\n            bbs_arrays = [np.array(x) for x in bbs_lists]\n            grid_BBs.append(bbs_arrays)\n\n        return np.array(grid_imgs), grid_BBs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"showcase_grid(\"afa5e8098\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"showcase_bbs(\"afa5e8098\", 43)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_gen = MyDataSequenceSavingBBs(3, TRAIN_IMG_IDS[:1])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batchx,batchy = train_gen[0]\ndraw_img_boxes(batchx[1],batchy[1])","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}