{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":53666,"databundleVersionId":6589269,"sourceType":"competition"},{"sourceId":154108503,"sourceType":"kernelVersion"}],"dockerImageVersionId":30587,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, json\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom matplotlib import pyplot as plt\nimport torch, torchvision\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\n# sampled 24 hours in integers\nTIME_WIDTH = 1024\n# start time for each window\nREF_TIME = '12:00:00.000'\n# amount of pixel to extend event bound box on both sides\nEVENT_EXTENSION = 21 # 21.333 ~ 30 mins\n# amount of pixel to extend sleep bound box on both sides\nSLEEP_EXTENSION = 85 # 85.333 ~ 2 hours\n\nDIR_TRAIN = '/kaggle/input/cmi-data-prep-v2/images'\nDIR_COOR = '/kaggle/input/cmi-data-prep-v2/coordinates'\nTRAIN_EVENTS_PATH = '/kaggle/input/child-mind-institute-detect-sleep-states/train_events.csv'\nMODEL_PATH = '/kaggle/working/model.pth'\nTEMP_MODEL_PATH = '/kaggle/working/temp_model.pth'\nINVALID_SIDS = ['1087d7b0ff2e', '44a41bba1ee7', '519ae2d858b0', '9aed9ee12ae2', 'c38707ef76df'] #'148471991ffb'\n# 1087d7b0ff2e: data not present in train_series\n# 44a41bba1ee7,519ae2d858b0,9aed9ee12ae2,c38707ef76df: wakeup too close to 12pm (11:59am)\ntorch_cache_path = 'torch_cache'\nif not os.path.exists(torch_cache_path):\n    os.makedirs(torch_cache_path)\ntorch.hub.set_dir(torch_cache_path)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-12-08T11:26:53.556295Z","iopub.execute_input":"2023-12-08T11:26:53.556690Z","iopub.status.idle":"2023-12-08T11:26:57.391661Z","shell.execute_reply.started":"2023-12-08T11:26:53.556659Z","shell.execute_reply":"2023-12-08T11:26:57.390843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AccelerometerDataset(Dataset):\n    def __init__(self, events_df: pd.DataFrame, images_dir: str, coordinates_dir: str, transforms=None) -> None:\n        super().__init__()\n\n        self.events_df = events_df\n        self.images_dir = images_dir\n        self.coordinates_dir = coordinates_dir\n        self.image_ids = self.get_image_ids(images_dir)\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n\n        image = cv2.imread(f'{self.images_dir}/{image_id}.png', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        ### Get onset and wakeup times ###\n        series_id, day_str = image_id.split('_')\n        night = int(day_str) + 1 # dataset uses 1-based indexing\n        records = self.events_df[(self.events_df['series_id']==series_id) & (self.events_df['night']==night)]\n\n        # Get start_timestamp offset\n        coor_path = os.path.join(self.coordinates_dir, series_id+'.json')\n        with open(coor_path) as f:\n            coor_json = json.load(f)\n        offset = float(coor_json[day_str]['offset'])\n        if offset == 6.0:\n            print(f'coor_path={coor_path} image_id={image_id}')\n\n        # Get bounds and labels\n        bounds = []\n        labels = [] # 0: background, 1: onset, 2: wakeup, 3: sleep\n        onset_idx = None\n        start_timestamp = pd.to_datetime(records['timestamp'].min().strftime('%Y-%m-%d') + ' ' + REF_TIME)\n        if records['timestamp'].min() < start_timestamp:\n            start_timestamp = start_timestamp - pd.to_timedelta(1, 'D')\n        start_timestamp = start_timestamp + pd.to_timedelta(offset, 'hours')\n        ref_epoch = int(start_timestamp.timestamp())\n        for idx in records.index:\n            event = records['event'][idx]\n            timestamp_epoch = int(records['timestamp'][idx].timestamp())\n            x_coor = round(TIME_WIDTH*(timestamp_epoch - ref_epoch)/86400) # 86400 seconds in a day\n            \n            if event == 'onset':\n                if onset_idx is None:\n                    onset_idx = [len(bounds)-1,idx]\n                else:\n                    # previous onset has no wakeup\n                    print(f'Missing wakeup detected: sid={series_id} night={night}')\n            else:\n                if onset_idx is None:\n                    # current wakeup does not follow an onset\n                    print(f'Missing onset detected: sid={series_id} night={night}')\n                else:\n                    # get sleep region bounding box\n                    x_min, x_max = self.get_x_min_max(bounds[onset_idx[0]][0], x_coor, SLEEP_EXTENSION, image.shape[1])\n                    y_min = 0\n                    y_max = image.shape[0]\n                    bounds.append([x_min, y_min, x_max, y_max])\n                    labels.append(3)\n                    onset_idx = None # reset\n                    \n            # get event region bounding box\n            x_min, x_max = self.get_x_min_max(x_coor, x_coor, EVENT_EXTENSION, image.shape[1])\n            y_min = 0\n            y_max = image.shape[0]\n            bounds.append([x_min, y_min, x_max, y_max])\n            if event == 'onset':\n                labels.append(1)\n            else:\n                labels.append(2)\n\n        assert len(labels) == len(bounds)\n        bounds = np.array(bounds)\n        labels = torch.as_tensor(labels, dtype=torch.int64)\n        area = (bounds[:, 3] - bounds[:, 1]) * (bounds[:, 2] - bounds[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        target = {}\n        target['boxes'] = bounds\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n        return image, target, image_id\n    \n    def __len__(self) -> int:\n        return len(self.image_ids)\n    \n    def get_image_ids(self, images_dir: str) -> list[str]:\n        f = []\n        for (dirpath, dirnames, filenames) in os.walk(images_dir):\n            f.extend(filenames)\n            break\n        f = [png_filename.split('.')[-2] for png_filename in f if png_filename.split('.')[-1]=='png']\n        f_annotated = []\n        for fname in f:\n            series_id, day_str = fname.split('_')\n            night = int(day_str) + 1\n            records = self.events_df[(self.events_df['series_id']==series_id) & (self.events_df['night']==night)]\n            #TODO remove those with only single or none annotations\n            if records.shape[0] != 0:\n                f_annotated.append(fname)\n        return f_annotated\n    \n    def get_x_min_max(self, x_left, x_right, extension, img_width):\n        x_min = max(0, x_left - extension)\n        x_max = min(img_width - 1, x_right + extension)\n        if (x_left - x_min) > (x_max - x_right):\n            x_min = x_left - (x_max - x_right)\n        elif (x_max - x_right) > (x_left - x_min):\n            x_max = x_right + (x_left - x_min)\n        return x_min, x_max\n    \n\nclass Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n\n# Albumentations\ndef get_train_transform_flip():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\ndef get_train_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n\ndef process_events_df(train_df: pd.DataFrame) -> pd.DataFrame:\n    train_df = train_df.dropna()\n    train_df = preproc_timestamp(train_df)\n    train_df = train_df.sort_values(by=['series_id','step'], ascending=[True, True])\n\n    # exclude out of range\n    rows_to_exclude = []\n    sids = train_df['series_id'].unique().tolist()\n    for sid in sids:\n        sid_df = train_df[train_df['series_id'] == sid]\n        nights = sid_df['night'].unique().tolist()\n        for night in nights:\n            night_df = sid_df[sid_df['night'] == night]\n            onset_row = night_df[night_df['event'] == 'onset']\n            if onset_row['timestamp'].shape[0] == 0:\n                print(f'skipping {sid} night={night}')\n                rows_to_exclude.append((sid,night))\n                continue\n            wakeup_row = night_df[night_df['event'] == 'wakeup']\n            if wakeup_row['timestamp'].shape[0] == 0:\n                print(f'skipping {sid} night={night}')\n                rows_to_exclude.append((sid,night))\n                continue\n\n    train_df = train_df[~train_df.set_index(['series_id', 'night']).index.isin(rows_to_exclude)]\n    train_df = train_df[~train_df.set_index(['series_id']).index.isin(INVALID_SIDS)]\n\n    return train_df\n\ndef preproc_timestamp(df):\n    df['date'] = df['timestamp'].str.split('T', expand=True)[0]\n    df['time'] = df['timestamp'].str.split('T', expand=True)[1].str.split('-', expand=True)[0]\n    df['timestamp'] = pd.to_datetime(df['date']+' '+df['time'])\n    df['hour'] = df['timestamp'].dt.hour\n    return df","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == '__main__':\n    #freeze_support()\n    train_df = pd.read_csv(TRAIN_EVENTS_PATH)\n    train_df = process_events_df(train_df)\n    series_ids = train_df['series_id'].unique().tolist()\n    split_point = round(0.2*len(series_ids))\n    valid_ids = series_ids[-split_point:]\n    train_ids = series_ids[:-split_point]\n    valid_df = train_df[train_df['series_id'].isin(valid_ids)]\n\n    # load a model; pre-trained on COCO\n    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, progress=True)\n\n    num_classes = 4  # 0: background, 1: onset, 2: wakeup, 3: sleep\n\n    # get number of input features for the classifier\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n\n    # replace the pre-trained head with a new one\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\n    train_dataset = AccelerometerDataset(train_df, DIR_TRAIN, DIR_COOR, get_train_transform())\n    valid_dataset = AccelerometerDataset(valid_df, DIR_TRAIN, DIR_COOR, get_valid_transform())\n\n    indices = torch.randperm(len(train_dataset)).tolist() #TODO split the dataset in train and test set\n    train_data_loader = DataLoader(\n        train_dataset,\n        batch_size=16,\n        shuffle=False,\n        num_workers=4,\n        collate_fn=collate_fn\n    )\n    valid_data_loader = DataLoader(\n        valid_dataset,\n        batch_size=8,\n        shuffle=False,\n        num_workers=4,\n        collate_fn=collate_fn\n    )\n    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\n    # Sample\n    images, targets, image_ids = next(iter(train_data_loader))\n    images = list(image.to(device) for image in images)\n    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n    boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)\n    sample = images[0].permute(1,2,0).cpu().numpy()\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    for box in boxes:\n        cv2.rectangle(sample,\n                    (box[0], box[1]),\n                    (box[2], box[3]),\n                    (220, 0, 0), 3)\n    #ax.set_axis_off()\n    ax.imshow(sample)\n    plt.show()\n    cv2.imwrite('temp1.png',sample)\n    print('temp1.png created')\n\n    ###### TRAINING #####\n    model.to(device)\n    params = [p for p in model.parameters() if p.requires_grad]\n    optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\n\n    loss_hist = Averager()\n    itr = 1\n    num_epochs = 50\n    for epoch in range(num_epochs):\n        loss_hist.reset()\n        for images, targets, image_ids in train_data_loader:\n            images = list(image.to(device) for image in images)\n            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n            loss_dict = model(images, targets)\n\n            losses = sum(loss for loss in loss_dict.values())\n            loss_value = losses.item()\n\n            loss_hist.send(loss_value)\n\n            optimizer.zero_grad()\n            losses.backward()\n            optimizer.step()\n\n            if itr % 50 == 0:\n                print(f\"Iteration #{itr} loss: {loss_value}\")\n\n            itr += 1\n        \n        # update the learning rate\n        if lr_scheduler is not None:\n            lr_scheduler.step()\n            \n        torch.save(model.state_dict(), TEMP_MODEL_PATH)\n\n        print(f\"Epoch #{epoch} loss: {loss_hist.value}\")\n        \n    torch.save(model.state_dict(), MODEL_PATH)\n\n    # Examine results\n    images, targets, image_ids = next(iter(valid_data_loader))\n    images = list(img.to(device) for img in images)\n    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n    boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)\n    sample = images[0].permute(1,2,0).cpu().numpy()\n    model.eval()\n    cpu_device = torch.device(\"cpu\")\n    outputs = model(images)\n    outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]\n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    for box in boxes:\n        cv2.rectangle(sample,\n                    (box[0], box[1]),\n                    (box[2], box[3]),\n                    (220, 0, 0), 3)\n    #ax.set_axis_off()\n    ax.imshow(sample)\n    plt.show()\n    cv2.imwrite('temp2.png',sample)\n    print('temp2.png created')\n\n    print('END')","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-12-08T10:28:42.433895Z","iopub.execute_input":"2023-12-08T10:28:42.434671Z","iopub.status.idle":"2023-12-08T10:30:01.647586Z","shell.execute_reply.started":"2023-12-08T10:28:42.434635Z","shell.execute_reply":"2023-12-08T10:30:01.645635Z"},"trusted":true},"execution_count":null,"outputs":[]}]}