{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install ultralytics -q\n!pip install scipy==1.8.0","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:52:22.080160Z","iopub.execute_input":"2023-08-06T10:52:22.080681Z","iopub.status.idle":"2023-08-06T10:53:03.822167Z","shell.execute_reply.started":"2023-08-06T10:52:22.080644Z","shell.execute_reply":"2023-08-06T10:53:03.821075Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport os\nfrom tqdm import tqdm\nimport wandb\nfrom ultralytics import YOLO\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom skimage.transform import resize\nfrom skimage import img_as_bool\nfrom skimage.morphology import convex_hull_image\nimport json","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-08-06T10:53:03.825755Z","iopub.execute_input":"2023-08-06T10:53:03.826066Z","iopub.status.idle":"2023-08-06T10:53:08.208072Z","shell.execute_reply.started":"2023-08-06T10:53:03.826037Z","shell.execute_reply":"2023-08-06T10:53:08.206954Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.init(mode='disabled')","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:08.209625Z","iopub.execute_input":"2023-08-06T10:53:08.210036Z","iopub.status.idle":"2023-08-06T10:53:09.288335Z","shell.execute_reply.started":"2023-08-06T10:53:08.209997Z","shell.execute_reply":"2023-08-06T10:53:09.287285Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Get rle encoding for given mask","metadata":{}},{"cell_type":"code","source":"def rle_encode(mask: np.ndarray) -> str:\n    # flatten the mask column wise\n    pixels = mask.T.flatten()\n\n    # pad the list of pixels to ensure they have leading and trailing 0\n    use_padding = False\n    if pixels[0] or pixels[-1]:\n        use_padding = True\n        pixel_padded = np.zeros([len(pixels) + 2], dtype=pixels.dtype)\n        pixel_padded[1:-1] = pixels\n        pixels = pixel_padded\n\n    # get the pixel indices where consecutive pixels don't match => binary masks begin or end\n    # +1 to convert to inclusive start indices, exclusive end indices of masks\n    # +1 to convert to 1-indexed pixel ordering\n    rle = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    if use_padding:\n        rle = rle - 1  # remove the extra offset introduced by padding\n\n    # end pixel indices of masks - start pixel indices of masks = length of RLE runs\n    # store the RLE run lengths replacing the end pixel indices \n    rle[1::2] = rle[1::2] - rle[:-1:2]\n\n    return ' '.join(str(x) for x in rle)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.291403Z","iopub.execute_input":"2023-08-06T10:53:09.292067Z","iopub.status.idle":"2023-08-06T10:53:09.301589Z","shell.execute_reply.started":"2023-08-06T10:53:09.292032Z","shell.execute_reply":"2023-08-06T10:53:09.300352Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Preprocessing","metadata":{}},{"cell_type":"code","source":"def preprocess(img_path, img, flags):\n    img2 = img\n    if flags['hist']:\n        img1 = histEq(img_path)\n        img2 = channel_shift(img, img1)\n    if flags['bz']:\n        img2 = binarize(img)\n    return img2","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.303148Z","iopub.execute_input":"2023-08-06T10:53:09.303764Z","iopub.status.idle":"2023-08-06T10:53:09.313139Z","shell.execute_reply.started":"2023-08-06T10:53:09.303731Z","shell.execute_reply":"2023-08-06T10:53:09.312124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def channel_shift(img, img1):\n    img2 = np.zeros_like(img)\n    img2[:,:,0] = img1\n    img2[:,:,1] = img1\n    img2[:,:,2] = img1\n    return img2","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.314461Z","iopub.execute_input":"2023-08-06T10:53:09.314860Z","iopub.status.idle":"2023-08-06T10:53:09.323079Z","shell.execute_reply.started":"2023-08-06T10:53:09.314828Z","shell.execute_reply":"2023-08-06T10:53:09.322183Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def binarize(img):\n    gim = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n    gim = cv2.adaptiveThreshold(gim, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 45, 11)\n    g3im = cv2.cvtColor(gim, cv2.COLOR_GRAY2BGR)\n    return g3im","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.324603Z","iopub.execute_input":"2023-08-06T10:53:09.324959Z","iopub.status.idle":"2023-08-06T10:53:09.334089Z","shell.execute_reply.started":"2023-08-06T10:53:09.324928Z","shell.execute_reply":"2023-08-06T10:53:09.333212Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def histEq(imgPath):\n    img= cv2.imread(imgPath,cv2.IMREAD_GRAYSCALE);\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))\n    cl1 = clahe.apply(img)\n    return cl1;","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.335537Z","iopub.execute_input":"2023-08-06T10:53:09.335940Z","iopub.status.idle":"2023-08-06T10:53:09.344057Z","shell.execute_reply.started":"2023-08-06T10:53:09.335907Z","shell.execute_reply":"2023-08-06T10:53:09.343163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Post Processing","metadata":{}},{"cell_type":"code","source":"def tableConvexHull(img, masks):\n    mask=np.zeros(masks[0].shape,dtype=\"bool\")\n    for msk in masks:\n        temp=msk.cpu().detach().numpy();\n        chull = convex_hull_image(temp);\n        mask=np.bitwise_or(mask,chull)\n    return mask","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.345461Z","iopub.execute_input":"2023-08-06T10:53:09.346121Z","iopub.status.idle":"2023-08-06T10:53:09.358664Z","shell.execute_reply.started":"2023-08-06T10:53:09.346090Z","shell.execute_reply":"2023-08-06T10:53:09.357715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def cls_exists(clss, cls):\n    indices = torch.where(clss==cls)\n    return len(indices[0])>0","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.362911Z","iopub.execute_input":"2023-08-06T10:53:09.363650Z","iopub.status.idle":"2023-08-06T10:53:09.369983Z","shell.execute_reply.started":"2023-08-06T10:53:09.363620Z","shell.execute_reply":"2023-08-06T10:53:09.368987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def empty_mask(img):\n    mask = np.zeros(img.shape[:2], dtype=\"uint8\")\n    return np.array(mask, dtype=bool)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.371220Z","iopub.execute_input":"2023-08-06T10:53:09.372262Z","iopub.status.idle":"2023-08-06T10:53:09.379279Z","shell.execute_reply.started":"2023-08-06T10:53:09.372230Z","shell.execute_reply":"2023-08-06T10:53:09.378510Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def extract_img_mask(img_model, img, config):\n    res_dict = {\n        'status' : 1\n    }\n    res = get_predictions(img_model, img, config)\n    \n    if res['status']==-1:\n        res_dict['status'] = -1\n        \n    elif res['status']==0:\n        res_dict['mask']=empty_mask(img)\n        \n    else:\n        masks = res['masks']\n        boxes = res['boxes']\n        clss = boxes[:, 5]\n        mask = extract_mask(img, masks, boxes, clss, 0)\n        res_dict['mask'] = mask\n    \n    return res_dict","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.380207Z","iopub.execute_input":"2023-08-06T10:53:09.383086Z","iopub.status.idle":"2023-08-06T10:53:09.390485Z","shell.execute_reply.started":"2023-08-06T10:53:09.383055Z","shell.execute_reply":"2023-08-06T10:53:09.389622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Getting predictions and extracting masks","metadata":{}},{"cell_type":"code","source":"def get_predictions(model, img2, config):\n    res_dict = {\n        'status': 1\n    }\n    try:\n        for result in model.predict(source=img2, verbose=False, retina_masks=config['rm'],\\\n                                    imgsz=config['sz'], conf=config['conf'], stream=True,\\\n                                    classes=config['classes']):\n            try:\n                res_dict['masks'] = result.masks.data\n                res_dict['boxes'] = result.boxes.data\n                del result\n                return res_dict\n            except Exception as e:\n                res_dict['status'] = 0\n                return res_dict\n    except:\n        res_dict['status'] = -1\n        return res_dict","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.391613Z","iopub.execute_input":"2023-08-06T10:53:09.392096Z","iopub.status.idle":"2023-08-06T10:53:09.401741Z","shell.execute_reply.started":"2023-08-06T10:53:09.392066Z","shell.execute_reply":"2023-08-06T10:53:09.400847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def extract_mask(img, masks, boxes, clss, cls):\n    if not cls_exists(clss, cls):\n        return empty_mask(img)\n    indices = torch.where(clss==cls)\n    c_masks = masks[indices]\n    mask_arr = torch.any(c_masks, dim=0).bool()\n    mask_arr = mask_arr.cpu().detach().numpy()\n    mask = mask_arr\n    return mask","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.403100Z","iopub.execute_input":"2023-08-06T10:53:09.403585Z","iopub.status.idle":"2023-08-06T10:53:09.412949Z","shell.execute_reply.started":"2023-08-06T10:53:09.403454Z","shell.execute_reply":"2023-08-06T10:53:09.411935Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_masks(filename, model, img_model, flags, configs):\n    response = {\n        'status': 1\n    }\n    ans_masks = []\n    img_path = os.path.join(img_dir, filename)\n    img = cv2.imread(img_path)\n    img2 = preprocess(img_path, img, flags)\n    \n    \n#     ***** Getting paragraph and text masks\n    res = get_predictions(model, img2, configs['paratext'])\n    if res['status']==-1:\n        response['status'] = -1\n        return response\n    elif res['status']==0:\n        for i in range(2): ans_masks.append(empty_mask(img))\n    else:\n        masks, boxes = res['masks'], res['boxes']\n        clss = boxes[:, 5]\n        for cls in range(2):\n            mask = extract_mask(img, masks, boxes, clss, cls)\n            ans_masks.append(mask)\n            \n            \n#     ***** Getting image and table masks\n    res2 = get_predictions(model, img2, configs['imgtab'])\n    if res2['status']==-1:\n        response['status'] = -1\n        return response\n    elif res2['status']==0:\n        for i in range(2): ans_masks.append(empty_mask(img))\n    else:\n        masks, boxes = res2['masks'], res2['boxes']\n        clss = boxes[:, 5]\n        \n        if cls_exists(clss, 2):\n            img_res = extract_img_mask(img_model, img, configs['image'])\n            if img_res['status'] == 1:\n                img_mask = img_res['mask']\n            else:\n                response['status'] = -1\n                return response\n            \n        else:\n            img_mask = empty_mask(img)\n        ans_masks.append(img_mask)\n        \n        if cls_exists(clss, 3):\n            indices = torch.where(clss==3)\n            tbl_mask = tableConvexHull(img, masks[indices])\n        else:\n            tbl_mask = empty_mask(img)\n        ans_masks.append(tbl_mask)\n    \n    if not configs['paratext']['rm']:\n        h, w, c = img.shape\n        for i in range(4):\n            ans_masks[i] = img_as_bool(resize(ans_masks[i], (h, w)))\n            \n    \n    response['masks'] = ans_masks\n    return response\n        ","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.414376Z","iopub.execute_input":"2023-08-06T10:53:09.414863Z","iopub.status.idle":"2023-08-06T10:53:09.429961Z","shell.execute_reply.started":"2023-08-06T10:53:09.414831Z","shell.execute_reply":"2023-08-06T10:53:09.429194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Putting it all together","metadata":{}},{"cell_type":"code","source":"def get_img_ids(lbl_file):\n    test_labels = open(lbl_file)\n    labels = json.loads(test_labels.read())\n    img_data = labels['images']\n    id_dict = {}\n    print('***** getting image ids *****')\n    for x in tqdm(img_data):\n        id_dict[x['file_name']] = x['id']\n    return id_dict","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.433140Z","iopub.execute_input":"2023-08-06T10:53:09.433641Z","iopub.status.idle":"2023-08-06T10:53:09.443958Z","shell.execute_reply.started":"2023-08-06T10:53:09.433610Z","shell.execute_reply":"2023-08-06T10:53:09.443117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def eval(df, files, model1, img_model, flags, configs, img_ids):\n    print('Starting evaluation on '+str(len(files))+\\\n          ' files with retina masks turned '+\\\n          ('ON' if configs['paratext']['rm'] else 'OFF'))\n    mo = 0\n    exc = []\n    for file in tqdm(files):\n        response = get_masks(file, model1, img_model, flags, configs)\n        if response['status']==-1:\n            exc.append(file)\n        else:\n            masks = response['masks']\n            img_id = img_ids[file]\n            for i, mask in enumerate(masks):\n                rle = rle_encode(mask)\n                df.loc[str(img_id)+'_'+str(i), 'Predicted'] = rle\n                \n    print('evaluation finished with retina masks turned '+\\\n         ('ON' if configs['paratext']['rm'] else 'OFF'))\n    \n    new_configs = configs.copy()\n    for key in new_configs.keys():\n#         new_configs[key] = configs[key].copy()\n        new_configs[key]['rm'] = False\n    \n    if len(exc)>0:\n        eval(df, exc, model1, img_model, flags, new_configs, img_ids)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.445195Z","iopub.execute_input":"2023-08-06T10:53:09.445581Z","iopub.status.idle":"2023-08-06T10:53:09.458817Z","shell.execute_reply.started":"2023-08-06T10:53:09.445549Z","shell.execute_reply":"2023-08-06T10:53:09.457963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Input Variables\nadjust these values to change weights and configurations","metadata":{}},{"cell_type":"code","source":"lbl_dir = '/kaggle/input/dlsprint2/badlad/badlad-test-metadata.json'\nimg_dir = '/kaggle/input/dlsprint2/badlad/images/test'\nimages = os.listdir(img_dir)\nid_dict = get_img_ids(lbl_dir)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.460267Z","iopub.execute_input":"2023-08-06T10:53:09.464061Z","iopub.status.idle":"2023-08-06T10:53:09.940816Z","shell.execute_reply.started":"2023-08-06T10:53:09.464030Z","shell.execute_reply":"2023-08-06T10:53:09.939933Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df1 = pd.read_csv('/kaggle/input/dlsprint2/sample_submission.csv')\ndf1 = df1.set_index('Id')","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:09.942024Z","iopub.execute_input":"2023-08-06T10:53:09.942450Z","iopub.status.idle":"2023-08-06T10:53:10.027922Z","shell.execute_reply.started":"2023-08-06T10:53:09.942417Z","shell.execute_reply":"2023-08-06T10:53:10.026985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"configs = {}\nconfigs['paratext'] = {\n    'sz' : 640,\n    'conf': 0.25,\n    'rm': True,\n    'classes': [0, 1]\n}\nconfigs['imgtab'] = {\n    'sz' : 640,\n    'conf': 0.35,\n    'rm': True,\n    'classes': [2, 3]\n}\nconfigs['image'] = {\n    'sz' : 640,\n    'conf': 0.35,\n    'rm': True,\n    'classes': [0]\n}\n\nflags = {\n    'hist' : False,\n    'bz' : False\n}","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:10.029311Z","iopub.execute_input":"2023-08-06T10:53:10.029635Z","iopub.status.idle":"2023-08-06T10:53:10.036397Z","shell.execute_reply.started":"2023-08-06T10:53:10.029604Z","shell.execute_reply":"2023-08-06T10:53:10.035516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_dir = '/kaggle/input/yolo-models/'\nmodel_name = \"e50_aug.pt\"\nimg_model_name = 'e100_img.pt'\nmodel1 = YOLO(base_dir+model_name)\nm_img = YOLO(base_dir+img_model_name)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:10.037745Z","iopub.execute_input":"2023-08-06T10:53:10.038304Z","iopub.status.idle":"2023-08-06T10:53:11.498172Z","shell.execute_reply.started":"2023-08-06T10:53:10.038273Z","shell.execute_reply":"2023-08-06T10:53:11.497187Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Running Inference","metadata":{}},{"cell_type":"code","source":"eval(df1, images, model1, m_img, flags, configs, id_dict)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:53:11.499695Z","iopub.execute_input":"2023-08-06T10:53:11.500086Z","iopub.status.idle":"2023-08-06T10:58:09.901012Z","shell.execute_reply.started":"2023-08-06T10:53:11.500054Z","shell.execute_reply":"2023-08-06T10:58:09.899238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(configs)","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:58:09.902133Z","iopub.status.idle":"2023-08-06T10:58:09.903158Z","shell.execute_reply.started":"2023-08-06T10:58:09.902911Z","shell.execute_reply":"2023-08-06T10:58:09.902941Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df1.to_csv('submissions.csv')","metadata":{"execution":{"iopub.status.busy":"2023-08-06T10:58:09.904756Z","iopub.status.idle":"2023-08-06T10:58:09.905268Z","shell.execute_reply.started":"2023-08-06T10:58:09.905028Z","shell.execute_reply":"2023-08-06T10:58:09.905050Z"},"trusted":true},"execution_count":null,"outputs":[]}]}