{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":85240,"databundleVersionId":9622164,"sourceType":"competition"}],"dockerImageVersionId":30805,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from collections import defaultdict\nfrom datetime import datetime\nfrom functools import partial  # NOQA: F401\nimport os\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nimport torch\nfrom torch.utils.data import Dataset, DataLoader  # NOQA: F401\nfrom torchmetrics.detection import MeanAveragePrecision\nfrom torchvision.io import read_image, ImageReadMode\nfrom torchvision.models.detection.fcos import fcos_resnet50_fpn, FCOSClassificationHead\nfrom torchvision.transforms.v2 import functional as trx\nfrom tqdm import tqdm\n\nfrom line_profiler import profile\n\nop = os.path\nORG_LABELS = {\n    \"aegypti\": 0,\n    \"albopictus\": 1,\n    \"anopheles\": 2,\n    \"culex\": 3,\n    \"culiseta\": 4,\n    \"japonicus/koreicus\": 5,\n}\nLABELS = {k: v + 1 for k, v in ORG_LABELS.items()}\nID2LABEL = {v: k for k, v in LABELS.items()}\nBAR_FORMAT = \"{l_bar}{bar}| {n_fmt}/{total_fmt} {rate_fmt}{postfix}\"\nMIN_SIZE, MAX_SIZE = 800, 1333","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-06T05:01:50.609035Z","iopub.execute_input":"2024-12-06T05:01:50.609737Z","iopub.status.idle":"2024-12-06T05:01:57.591192Z","shell.execute_reply.started":"2024-12-06T05:01:50.609703Z","shell.execute_reply":"2024-12-06T05:01:57.590452Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TrainDataset(Dataset):\n    def __init__(self, df, root, device=None):\n        self.df = df\n        self.root = root\n        if device is None:\n            device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.device = device\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        path = op.join(self.root, \"images\", f\"{row['id']}.jpeg\")\n        image = read_image(path, ImageReadMode.RGB)\n        image = trx.to_dtype(image, torch.float32, scale=True)\n        _, im_height, im_width = image.shape\n        x_center, y_center, width, height = row[\n            [\"x_center\", \"y_center\", \"width\", \"height\"]\n        ]\n        x_center, width = x_center * im_width, width * im_width\n        y_center, height = y_center * im_height, height * im_height\n        x0, y0 = x_center - width / 2, y_center - height / 2\n        x1, y1 = x_center + width / 2, y_center + height / 2\n        boxes = torch.tensor([x0, y0, x1, y1], dtype=torch.float32).unsqueeze(0)\n        labels = torch.tensor([row[\"label\"] + 1], dtype=torch.int64)\n        return image, {\"boxes\": boxes, \"labels\": labels}\n\n    def show(self, n, model=None, seed=None):\n        fig, ax = plt.subplots(nrows=n, ncols=n, figsize=(3 * n, 3 * n))\n        if seed is not None:\n            np.random.seed(seed)\n        samples = [self[i] for i in np.random.choice(len(self), size=n * n)]\n        preds = []\n        for (im, box), ax in zip(samples, ax.ravel()):\n            ax.imshow(im.permute(1, 2, 0))\n            x0, y0, x1, y1 = box[\"boxes\"][0]\n            ax.add_patch(\n                plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=\"green\")\n            )\n            fig_label = ID2LABEL[box[\"labels\"][0].item()]\n            if model is not None:\n                model.eval()\n                with torch.no_grad():\n                    pred = model([im.to(self.device)])\n                preds.append(pred[0])\n                ix = pred[0][\"scores\"].argmax().item()\n                predbox = pred[0][\"boxes\"][ix]\n                predlabel = ID2LABEL[pred[0][\"labels\"][ix].item()]\n                x0, y0, x1, y1 = predbox.cpu().numpy()\n                ax.add_patch(\n                    plt.Rectangle(\n                        (x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=\"red\"\n                    )\n                )\n                fig_label = f\"{fig_label} ({predlabel})\"\n            ax.set_axis_off()\n            ax.set_title(fig_label, fontsize=\"small\")\n        plt.tight_layout()\n        plt.savefig(\"foo.png\", bbox_inches=\"tight\")\n        return preds, samples\n\n\ndef make_df(root, stratify=False, test_size=0.2):\n    images_path = op.join(root, \"images\")\n    labels_path = op.join(root, \"labels\")\n    image_files = [\n        op.join(images_path, f) for f in os.listdir(images_path) if f.endswith(\".jpeg\")\n    ]\n    payload = []\n    for file in image_files:\n        image_id = op.splitext(op.basename(file))[0]\n        label_file = op.join(labels_path, f\"{image_id}.txt\")\n        with open(label_file, \"r\") as f_in:\n            label, x_center, y_center, width, height = map(\n                float, f_in.read().strip().split()\n            )\n            label = int(label)\n\n        payload.append(\n            dict(\n                id=image_id,\n                label=label,\n                x_center=x_center,\n                y_center=y_center,\n                width=width,\n                height=height,\n            )\n        )\n    df = pd.DataFrame.from_records(payload)\n    if not stratify:\n        return df\n    area = df[\"width\"] * df[\"height\"]\n    area_cat = pd.cut(area, bins=np.linspace(0, 1, 6), labels=list(\"ABCDE\"))\n    strf_cat = df[\"label\"].astype(int).astype(str) + area_cat.astype(str)\n    return train_test_split(df, stratify=strf_cat, test_size=test_size)\n\n\ndef get_model(device=None):\n    model = fcos_resnet50_fpn(weights=\"COCO_V1\")  # , trainable_backbone_layers=0)\n    for p in model.parameters():\n        p.requires_grad = False\n    model.head.classification_head = FCOSClassificationHead(\n        in_channels=256, num_anchors=model.head.classification_head.num_anchors,\n        num_classes=len(LABELS)+1\n    )\n    if device is None:\n        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    return model.to(device)\n\n\ndef load_model(path):\n    model = get_model()\n    model.load_state_dict(torch.load(path, weights_only=True))\n    return model\n\n\ndef resize_with_ar(image, box):\n    \"\"\"Resize an image to exactly MIN_SIZE by MAX_SIZE. Transpose if needed.\"\"\"\n    h, w = image.shape[-2:]\n    if h < w:  # landscape\n        new_image = trx.resize(image, (MIN_SIZE, MAX_SIZE))\n        box[::2] *= MAX_SIZE / w\n        box[1::2] *= MIN_SIZE / h\n    else:\n        new_image = trx.resize(image, (MAX_SIZE, MIN_SIZE)).permute(0, 2, 1)\n        box[::2] *= MIN_SIZE / w\n        box[1::2] *= MAX_SIZE / h\n    return new_image, box\n\n\ndef collate(batch, device=False, resize=True):\n    images, targets = zip(*batch)\n    if resize:\n        images, bboxes = zip(*[resize_with_ar(im, t[\"boxes\"][0]) for im, t in zip(images, targets)])\n        new_targets = []\n        for b, l in zip(bboxes, targets):\n            l[\"boxes\"] = b.unsqueeze(0)\n            l['labels'] = l['labels']\n            new_targets.append(l)\n        images = torch.stack(images)\n        if device:\n            images = images.to(device)\n            new_targets = [{'labels': i['labels'].to(device), 'boxes': i['boxes'].to(device)} for i in new_targets]\n    elif device:\n        images = [i.to(device) for i in images]\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n    return images, targets\n\n\ndef update_history(cache, patience=1, **kwargs):\n    report = []\n    for metric in [\"train_loss\", \"val_loss\", \"train_metric\", \"val_metric\"]:\n        if metric not in kwargs:\n            continue\n        cache[metric].append(kwargs[metric])\n        window = cache[metric][-patience:]\n        if len(cache[metric]) >= patience + 1:\n            direction = \"↓\" if np.mean(window) <= cache[metric][-patience - 1] else \"↑\"\n        else:\n            direction = \"⌾\"\n\n        # Last boolean denotes whether it is minimizing\n        report.append((metric, direction, window[-1], metric.endswith(\"_loss\")))\n    return report\n\n\ndef update_postfix(bar, report, lr):\n    bar.update(1)\n    postfix_params = {}\n    for metric, direction, value, minimizing in report:\n        value = round(value, 3)\n        if direction == \"⌾\":\n            colored = f\"\\033[33m{value}{direction}\\033[0m\"\n        if (minimizing and direction == \"↓\") or (not minimizing and direction == \"↑\"):\n            colored = f\"\\033[32m{value}{direction}\\033[0m\"\n        elif (minimizing and direction == \"↑\") or (not minimizing and direction == \"↓\"):\n            colored = f\"\\033[31m{value}{direction}\\033[0m\"\n        postfix_params[metric] = colored\n    postfix_params[\"lr\"] = f\"{round(lr[-1], 6)}\"\n    bar.set_postfix(postfix_params, refresh=True)\n\n\ndef train(model, train_loader, test_loader, n_epochs=1, patience=3):\n    opt = torch.optim.Adam(model.head.classification_head.parameters(), lr=0.001)\n    # lr = torch.optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.1)\n    b_len = len(train_loader) + len(test_loader)\n    epoch_history = defaultdict(list)  # NOQA: F841\n    with tqdm(desc=\"Epoch\", total=n_epochs, bar_format=BAR_FORMAT) as ebar:\n        for epoch in range(n_epochs):\n            model.train()\n            epoch_train_loss = epoch_val_loss = 0\n            epoch_val_metric = 0\n            for images, targets in train_loader:\n                targets = [{\n                    'labels': t['labels'].to(device),\n                    'boxes': t['boxes'].to(device)\n                } for t in targets]\n                opt.zero_grad()\n                loss_dict = model(images.to(device), targets)\n                batch_train_loss = sum(loss for loss in loss_dict.values())\n                batch_train_loss.backward()\n                btl = batch_train_loss.item()\n                opt.step()\n                epoch_train_loss += btl\n            epoch_train_loss /= len(train_loader)\n            for images, targets in test_loader:\n                targets = [{\n                    'labels': t['labels'].to(device),\n                    'boxes': t['boxes'].to(device)\n                } for t in targets]\n                with torch.no_grad():\n                    preds = model(images.to(device), targets)\n                batch_val_loss = sum(loss for loss in preds.values())\n                epoch_val_loss += batch_val_loss.item()\n                epoch_val_loss /= len(test_loader)\n                epoch_val_metric /= len(test_loader)\n            # lr.step()\n            report = update_history(\n                epoch_history,\n                train_loss=epoch_train_loss,\n                val_loss=epoch_val_loss,\n                val_metric=epoch_val_metric,\n                patience=patience,\n            )\n            update_postfix(ebar, report, [0.001])\n    return epoch_history\n\n\ndef make_submission(\n    model,\n    root,\n    outpath=None,\n    show=0,\n    seed=None,\n    dfpath=\"data/sample_submission.csv\",\n    device=None,\n):\n    if device is None:\n        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model.eval()\n    df = pd.read_csv(dfpath, index_col=\"ImageID\")\n    for idx, row in tqdm(df.iterrows(), total=len(df)):\n        path = op.join(root, idx)\n        image = read_image(path, ImageReadMode.RGB)\n        image = trx.to_dtype(image, torch.float32, scale=True)\n        with torch.no_grad():\n            pred = model([image.to(device)])[0]\n        try:\n            ix = pred[\"scores\"].argmax().item()\n            conf = pred[\"scores\"][ix].item()\n            predbox = pred[\"boxes\"][ix]\n            label = ID2LABEL[pred[\"labels\"][ix].item()]\n            x0, y0, x1, y1 = predbox.cpu().numpy()\n            _, im_height, im_width = image.shape\n            xcenter, ycenter = (x0 + x1) / 2, (y0 + y1) / 2\n            box_width, box_height = x1 - x0, y1 - y0\n            xcenter, box_width = xcenter / im_width, box_width / im_width\n            ycenter, box_height = ycenter / im_height, box_height / im_height\n        except IndexError:\n            label, conf, xcenter, ycenter, box_width, box_height = \"albopictus\", 0, 0, 0, 0, 0\n        df.loc[idx, ['LabelName', 'Conf', 'xcenter', 'ycenter', 'bbx_width', 'bbx_height']] = [\n            label, conf, xcenter, ycenter, box_width, box_height\n        ]\n    if outpath is None:\n        outpath = \"submission.csv\"\n    df.to_csv(outpath)\n    if show:\n        np.random.seed(seed)\n        xdf = df.loc[np.random.choice(df.index, size=(show * show))]\n        fig, ax = plt.subplots(nrows=show, ncols=show, figsize=(3 * show, 3 * show))\n        for (idx, row), ax in zip(xdf.iterrows(), ax.ravel()):\n            path = op.join(root, idx)\n            image = read_image(path, ImageReadMode.RGB)\n            _, im_height, im_width = image.shape\n            ax.imshow(image.permute(1, 2, 0))\n            xcenter, ycenter, box_width, box_height = row[[\"xcenter\", \"ycenter\", \"bbx_width\", \"bbx_height\"]]\n            x0, y0 = xcenter - box_width / 2, ycenter - box_height / 2\n            x1, y1 = xcenter + box_width / 2, ycenter + box_height / 2\n            x0, x1 = x0 * im_width, x1 * im_width\n            y0, y1 = y0 * im_height, y1 * im_height\n            ax.add_patch(\n                plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=\"red\")\n            )\n            ax.set_axis_off()\n            ax.set_title(f\"{row['LabelName']} ({round(row['Conf'], 2)})\", fontsize=\"small\")\n        plt.tight_layout()\n        plt.savefig(\"predictions.png\", bbox_inches=\"tight\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-06T07:13:03.951850Z","iopub.execute_input":"2024-12-06T07:13:03.952367Z","iopub.status.idle":"2024-12-06T07:13:04.458583Z","shell.execute_reply.started":"2024-12-06T07:13:03.952316Z","shell.execute_reply":"2024-12-06T07:13:04.457070Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nroot = \"/kaggle/input/dlp-object-detection/final_dlp_data/final_dlp_data/train/\"\ndftrain, dftest = make_df(root, stratify=True)\ndstrain, dstest = TrainDataset(dftrain, root), TrainDataset(dftest, root)\nmodel = get_model(device)\ntrain_loader = DataLoader(\n    dstrain, batch_size=16, shuffle=True,  num_workers=4,\n    collate_fn=partial(collate, resize=True),\n)\ntest_loader = DataLoader(\n    dstest, batch_size=8, shuffle=False, num_workers=2,\n    collate_fn=partial(collate, resize=True)\n\n)\nhistory = train(model, train_loader, test_loader, n_epochs=20)\nwith open(\"history.json\", \"w\") as f_out:\n    json.dump(history, f_out, indent=2)\nmake_submission(model, \"/kaggle/input/dlp-object-detection/final_dlp_data/final_dlp_data/test/images/\",\n                show=4, dfpath=\"/kaggle/input/dlp-object-detection/sample_submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-06T05:13:23.525752Z","iopub.execute_input":"2024-12-06T05:13:23.526564Z","iopub.status.idle":"2024-12-06T05:28:01.083086Z","shell.execute_reply.started":"2024-12-06T05:13:23.526529Z","shell.execute_reply":"2024-12-06T05:28:01.081735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}