{"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":"from PIL import Image\nimport numpy as np\nfrom matplotlib import patches\nfrom tqdm import tqdm\nimport pandas as pd\nimport fnmatch\nfrom multiprocessing import Pool\nfrom pathlib import Path\nfrom functools import partial","metadata":{"execution":{"iopub.status.busy":"2022-03-07T23:24:50.710956Z","iopub.execute_input":"2022-03-07T23:24:50.711699Z","iopub.status.idle":"2022-03-07T23:24:50.737161Z","shell.execute_reply.started":"2022-03-07T23:24:50.711607Z","shell.execute_reply":"2022-03-07T23:24:50.736538Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!git clone https://github.com/YangtaoWANG95/TokenCut.git","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:44.311299Z","iopub.execute_input":"2022-03-07T12:48:44.311573Z","iopub.status.idle":"2022-03-07T12:48:45.095622Z","shell.execute_reply.started":"2022-03-07T12:48:44.311537Z","shell.execute_reply":"2022-03-07T12:48:45.094726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cd TokenCut && pip install -r requirements.txt","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:45.096949Z","iopub.execute_input":"2022-03-07T12:48:45.097315Z","iopub.status.idle":"2022-03-07T12:48:53.098743Z","shell.execute_reply.started":"2022-03-07T12:48:45.097272Z","shell.execute_reply":"2022-03-07T12:48:53.097862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_img_folder = '../input/jpeg-happywhale-384x384/train_images-384-384/train_images-384-384/'\ntest_img_folder = '../input/jpeg-happywhale-384x384/test_images-384-384/test_images-384-384/'","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:53.102804Z","iopub.execute_input":"2022-03-07T12:48:53.103027Z","iopub.status.idle":"2022-03-07T12:48:53.109127Z","shell.execute_reply.started":"2022-03-07T12:48:53.102999Z","shell.execute_reply":"2022-03-07T12:48:53.108343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.insert(0, './TokenCut')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:53.110529Z","iopub.execute_input":"2022-03-07T12:48:53.1108Z","iopub.status.idle":"2022-03-07T12:48:53.117967Z","shell.execute_reply.started":"2022-03-07T12:48:53.110765Z","shell.execute_reply":"2022-03-07T12:48:53.11726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport argparse\nimport random\nimport pickle\n\nimport torch\nimport datetime\nimport torch.nn as nn\nimport numpy as np\n\nfrom tqdm import tqdm\nfrom PIL import Image\n\nfrom networks import get_model\nfrom datasets import ImageDataset, Dataset, bbox_iou\nfrom visualizations import visualize_img, visualize_eigvec, visualize_predictions, visualize_predictions_gt \nfrom object_discovery import ncut \nimport matplotlib.pyplot as plt\nimport time\n\n# torch.multiprocessing.set_start_method('spawn')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:53.119389Z","iopub.execute_input":"2022-03-07T12:48:53.119642Z","iopub.status.idle":"2022-03-07T12:48:54.856781Z","shell.execute_reply.started":"2022-03-07T12:48:53.119609Z","shell.execute_reply":"2022-03-07T12:48:54.856011Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATCH_SIZE = 16\nWHICH_FEATURES = 'k'\nARCH = 'vit_base'\n\nTAU = 0.2\nEPS = 1e-5\nNO_BINARY_GRAPH = False","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:54.858314Z","iopub.execute_input":"2022-03-07T12:48:54.858584Z","iopub.status.idle":"2022-03-07T12:48:54.865263Z","shell.execute_reply.started":"2022-03-07T12:48:54.858548Z","shell.execute_reply":"2022-03-07T12:48:54.864508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"exp_name = f\"TokenCut-{ARCH}\"\nif \"vit\" in ARCH:\n    exp_name += f\"{PATCH_SIZE}_{WHICH_FEATURES}\"","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:54.868215Z","iopub.execute_input":"2022-03-07T12:48:54.868413Z","iopub.status.idle":"2022-03-07T12:48:54.874055Z","shell.execute_reply.started":"2022-03-07T12:48:54.868387Z","shell.execute_reply":"2022-03-07T12:48:54.873268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_bounding_box(img_path, model):\n    dataset = ImageDataset(img_path)\n    preds_dict = {}\n    cnt = 0\n    corloc = np.zeros(len(dataset.dataloader))\n\n    start_time = time.time() \n    pbar = dataset.dataloader\n    for im_id, inp in enumerate(pbar):\n\n        # ------------ IMAGE PROCESSING -------------------------------------------\n        img = inp[0]\n\n        init_image_size = img.shape\n\n        # Get the name of the image\n        im_name = dataset.get_image_name(inp[1])\n        # Pass in case of no gt boxes in the image\n        if im_name is None:\n            continue\n\n        # Padding the image with zeros to fit multiple of patch-size\n        size_im = (\n            img.shape[0],\n            int(np.ceil(img.shape[1] / PATCH_SIZE) * PATCH_SIZE),\n            int(np.ceil(img.shape[2] / PATCH_SIZE) * PATCH_SIZE),\n        )\n        paded = torch.zeros(size_im)\n        paded[:, : img.shape[1], : img.shape[2]] = img\n        img = paded\n\n        # # Move to gpu\n        img = img.cuda(non_blocking=True)\n\n        # Size for transformers\n        w_featmap = img.shape[-2] // PATCH_SIZE\n        h_featmap = img.shape[-1] // PATCH_SIZE\n\n        # ------------ EXTRACT FEATURES -------------------------------------------\n        with torch.no_grad():\n\n            # ------------ FORWARD PASS -------------------------------------------\n            if \"vit\"  in ARCH:\n                # Store the outputs of qkv layer from the last attention layer\n                feat_out = {}\n                def hook_fn_forward_qkv(module, input, output):\n                    feat_out[\"qkv\"] = output\n                model._modules[\"blocks\"][-1]._modules[\"attn\"]._modules[\"qkv\"].register_forward_hook(hook_fn_forward_qkv)\n\n                # Forward pass in the model\n                attentions = model.get_last_selfattention(img[None, :, :, :])\n\n                # Scaling factor\n                scales = [PATCH_SIZE, PATCH_SIZE]\n\n                # Dimensions\n                nb_im = attentions.shape[0]  # Batch size\n                nh = attentions.shape[1]  # Number of heads\n                nb_tokens = attentions.shape[2]  # Number of tokens\n\n                # Extract the qkv features of the last attention layer\n                qkv = (\n                    feat_out[\"qkv\"]\n                    .reshape(nb_im, nb_tokens, 3, nh, -1 // nh)\n                    .permute(2, 0, 3, 1, 4)\n                )\n                q, k, v = qkv[0], qkv[1], qkv[2]\n                k = k.transpose(1, 2).reshape(nb_im, nb_tokens, -1)\n                q = q.transpose(1, 2).reshape(nb_im, nb_tokens, -1)\n                v = v.transpose(1, 2).reshape(nb_im, nb_tokens, -1)\n\n                # Modality selection\n                if WHICH_FEATURES == \"k\":\n                    #feats = k[:, 1:, :]\n                    feats = k\n                elif WHICH_FEATURES == \"q\":\n                    #feats = q[:, 1:, :]\n                    feats = q\n                elif WHICH_FEATURES == \"v\":\n                    #feats = v[:, 1:, :]\n                    feats = v\n\n\n\n            else:\n                raise ValueError(\"Unknown model.\")\n\n        # ------------ Apply TokenCut ------------------------------------------- \n        pred, objects, foreground, seed , bins, eigenvector= ncut(\n            feats, [w_featmap, h_featmap], scales, init_image_size, TAU, EPS, im_name=im_name, no_binary_graph=NO_BINARY_GRAPH)\n        return img, pred","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:54.87644Z","iopub.execute_input":"2022-03-07T12:48:54.876714Z","iopub.status.idle":"2022-03-07T12:48:54.893923Z","shell.execute_reply.started":"2022-03-07T12:48:54.876679Z","shell.execute_reply":"2022-03-07T12:48:54.89314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def convert_to_rect(img, label, color='b'):\n    height, width = img.shape[0], img.shape[1]\n    xmin, ymin, xmax, ymax  = label[0], label[1], label[2], label[3]\n    rect = patches.Rectangle((\n         xmin,\n         ymin\n    ),\n        (xmax - xmin),\n        (ymax - ymin),\n        linewidth=1, edgecolor=color, facecolor='none'\n    )\n    return rect","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:55.931919Z","iopub.execute_input":"2022-03-07T12:48:55.932622Z","iopub.status.idle":"2022-03-07T12:48:55.938878Z","shell.execute_reply.started":"2022-03-07T12:48:55.932575Z","shell.execute_reply":"2022-03-07T12:48:55.937992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _get_bb(path):\n    img, pred = get_bounding_box(str(path))\n    return path.name, pred","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:48:56.30774Z","iopub.execute_input":"2022-03-07T12:48:56.308611Z","iopub.status.idle":"2022-03-07T12:48:56.314037Z","shell.execute_reply.started":"2022-03-07T12:48:56.308573Z","shell.execute_reply":"2022-03-07T12:48:56.313121Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_model(arch, patch_size):\n    device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n    #device = torch.device('cuda')\n    model = get_model(arch, patch_size, device)\n        \n    return model","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:02.682234Z","iopub.execute_input":"2022-03-07T12:49:02.682523Z","iopub.status.idle":"2022-03-07T12:49:02.687417Z","shell.execute_reply.started":"2022-03-07T12:49:02.682488Z","shell.execute_reply":"2022-03-07T12:49:02.686414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_preds(model, img_folder, n=None):\n    total_len = len(fnmatch.filter(os.listdir(img_folder), '*.jpg'))\n    preds = []\n    pathlist = Path(img_folder).glob('*jpg')\n    for i, path in tqdm(enumerate(pathlist), total=n or total_len):\n        img, out_pred = get_bounding_box(str(path), model)\n        width, height = float(img.shape[1]), float(img.shape[2]) \n        xmin, ymin, xmax, ymax = list(out_pred)\n        preds.append([path.name, (xmin / width), (ymin / height), (xmax / width), (ymax / height)])\n        if n and i >= n:\n            break\n\n    return preds","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:04.224957Z","iopub.execute_input":"2022-03-07T12:49:04.225688Z","iopub.status.idle":"2022-03-07T12:49:04.233906Z","shell.execute_reply.started":"2022-03-07T12:49:04.225631Z","shell.execute_reply":"2022-03-07T12:49:04.233136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_preds(preds):\n    return pd.DataFrame(preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax'])","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:05.639861Z","iopub.execute_input":"2022-03-07T12:49:05.640646Z","iopub.status.idle":"2022-03-07T12:49:05.645699Z","shell.execute_reply.started":"2022-03-07T12:49:05.6406Z","shell.execute_reply":"2022-03-07T12:49:05.644876Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def convert_to_rect(img, label, linewidth=1, color='b'):\n    width, height = img.size[0], img.size[1]\n    xmin, ymin, xmax, ymax  = label[0], label[1], label[2], label[3]\n    rect = patches.Rectangle((\n         xmin * width,\n         ymin * height\n    ),\n        (xmax - xmin) * width,\n        (ymax - ymin) * height,\n        linewidth=linewidth, edgecolor=color, facecolor='none'\n    )\n    return rect","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:06.160026Z","iopub.execute_input":"2022-03-07T12:49:06.16094Z","iopub.status.idle":"2022-03-07T12:49:06.167804Z","shell.execute_reply.started":"2022-03-07T12:49:06.160889Z","shell.execute_reply":"2022-03-07T12:49:06.166883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"color_list = ['r', 'g', 'b']\n\ndef show_img_grid(preds, dataset):\n    row = 10; col = 4;\n\n    plt.figure(figsize=(20,int(20*row/col)))\n    for j in range(row*col):\n        first_preds = preds[0]\n\n        first_pred = first_preds[j]\n        image, bb = first_pred[0], first_pred[1:]\n        img = Image.open(f'../input/happy-whale-and-dolphin/{dataset}_images/{image}')\n        plt.subplot(row,col,j+1)\n        plt.axis('off')\n        plt.imshow(img)\n        ax = plt.gca()\n        \n        for i, pred in enumerate(preds):\n            bb = pred[j][1:]\n            c = color_list[i]\n            ax.add_patch(convert_to_rect(img, bb, 3, color=c))\n    plt.tight_layout()\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:07.87797Z","iopub.execute_input":"2022-03-07T12:49:07.878241Z","iopub.status.idle":"2022-03-07T12:49:07.885849Z","shell.execute_reply.started":"2022-03-07T12:49:07.87821Z","shell.execute_reply":"2022-03-07T12:49:07.885134Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = load_model('vit_small', 16)\npreds_small = get_preds(model, train_img_folder, n=100)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:14.873403Z","iopub.execute_input":"2022-03-07T12:49:14.874107Z","iopub.status.idle":"2022-03-07T12:49:23.44609Z","shell.execute_reply.started":"2022-03-07T12:49:14.874069Z","shell.execute_reply":"2022-03-07T12:49:23.445346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = load_model('vit_base', 16)\npreds_base = get_preds(model, train_img_folder, n=100)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:23.447905Z","iopub.execute_input":"2022-03-07T12:49:23.448543Z","iopub.status.idle":"2022-03-07T12:49:31.741734Z","shell.execute_reply.started":"2022-03-07T12:49:23.448503Z","shell.execute_reply":"2022-03-07T12:49:31.740944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Average performance.\n# model = load_model('moco_vit_small', 16)\n# preds_moco_vit_small = get_preds(model, train_img_folder, n=100)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:31.743128Z","iopub.execute_input":"2022-03-07T12:49:31.743577Z","iopub.status.idle":"2022-03-07T12:49:31.747177Z","shell.execute_reply.started":"2022-03-07T12:49:31.743537Z","shell.execute_reply":"2022-03-07T12:49:31.746394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = load_model('moco_vit_base', 16)\npreds_moco_vit_base = get_preds(model, train_img_folder, n=100)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:31.749284Z","iopub.execute_input":"2022-03-07T12:49:31.749737Z","iopub.status.idle":"2022-03-07T12:49:40.301627Z","shell.execute_reply.started":"2022-03-07T12:49:31.749701Z","shell.execute_reply":"2022-03-07T12:49:40.300814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model = load_model('mae_vit_base', 16)\n# preds_mae_vit_base = get_preds(model, train_img_folder, n=100)\n# Not good!","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:40.303047Z","iopub.execute_input":"2022-03-07T12:49:40.303481Z","iopub.status.idle":"2022-03-07T12:49:40.307045Z","shell.execute_reply.started":"2022-03-07T12:49:40.303425Z","shell.execute_reply":"2022-03-07T12:49:40.306309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_img_grid([preds_small, preds_base, preds_moco_vit_base], dataset='train')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:49:40.30843Z","iopub.execute_input":"2022-03-07T12:49:40.30892Z","iopub.status.idle":"2022-03-07T12:50:03.24406Z","shell.execute_reply.started":"2022-03-07T12:49:40.308879Z","shell.execute_reply":"2022-03-07T12:50:03.242556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"These bounding boxes seem basically flawless. \n\nLet's try on test.","metadata":{}},{"cell_type":"code","source":"model = load_model('vit_small', 16)\npreds_small = get_preds(model, test_img_folder, n=100)\n\nmodel = load_model('vit_base', 16)\npreds_base = get_preds(model, test_img_folder, n=100)\n\nmodel = load_model('moco_vit_base', 16)\npreds_moco_vit_base = get_preds(model, test_img_folder, n=100)\n\nshow_img_grid([preds_small, preds_base, preds_moco_vit_base], dataset='test')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:50:25.236248Z","iopub.execute_input":"2022-03-07T12:50:25.236747Z","iopub.status.idle":"2022-03-07T12:51:09.230815Z","shell.execute_reply.started":"2022-03-07T12:50:25.236703Z","shell.execute_reply":"2022-03-07T12:51:09.226578Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = load_model('vit_small', 16)\npreds = get_preds(model, train_img_folder)\npd.DataFrame(\n    preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax']\n).to_csv('train_vit_small.csv')\n\nmodel = load_model('vit_base', 16)\npreds = get_preds(model, train_img_folder)\npd.DataFrame(\n    preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax']\n).to_csv('train_vit_base.csv')\n\nmodel = load_model('moco_vit_base', 16)\npreds = get_preds(model, train_img_folder)\npd.DataFrame(\n    preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax']\n).to_csv('train_moco_vit_base.csv')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:52:06.645864Z","iopub.execute_input":"2022-03-07T12:52:06.646742Z","iopub.status.idle":"2022-03-07T12:52:10.182049Z","shell.execute_reply.started":"2022-03-07T12:52:06.646692Z","shell.execute_reply":"2022-03-07T12:52:10.180616Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = load_model('vit_small', 16)\ntest_preds = get_preds(model, test_img_folder)\npd.DataFrame(\n    test_preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax']\n).to_csv('test_vit_small.csv')\n\nmodel = load_model('vit_base', 16)\ntest_preds = get_preds(model, test_img_folder)\npd.DataFrame(\n    test_preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax']\n).to_csv('test_vit_base.csv')\n\nmodel = load_model('moco_vit_base', 16)\ntest_preds = get_preds(model, test_img_folder)\npd.DataFrame(\n    test_preds, columns=['image', 'xmin', 'ymin', 'xmax', 'ymax']\n).to_csv('test_moco_vit_base.csv')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls -l ./","metadata":{"execution":{"iopub.status.busy":"2022-03-07T12:52:12.447135Z","iopub.execute_input":"2022-03-07T12:52:12.447421Z","iopub.status.idle":"2022-03-07T12:52:13.242477Z","shell.execute_reply.started":"2022-03-07T12:52:12.447391Z","shell.execute_reply":"2022-03-07T12:52:13.241525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}