{"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":91531,"databundleVersionId":11044547,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install lightning lightning-bolts hydra-core natsort","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:23:30.088237Z","iopub.execute_input":"2025-03-13T05:23:30.088554Z","iopub.status.idle":"2025-03-13T05:23:39.150719Z","shell.execute_reply.started":"2025-03-13T05:23:30.088520Z","shell.execute_reply":"2025-03-13T05:23:39.149408Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config = {\n    \"model\": {\n        \"channels\": [2, 8,16, 32,48, 64, 80, 96, 112, 128, 256],  # Input channels and intermediate layers\n        \"t_kernel_size\": 7,  # Temporal kernel size\n        \"n_depthwise_layers\": 5,  # Depthwise layers for efficiency\n        \"detector_head\": True,  # Enable detector head\n        \"detector_depthwise\": True,  # Standard convolutions for detector\n        \"full_conv3d\": True,  # 1+2D convs instead of full 3D\n        \"norms\": \"all_bn\"  # BatchNorm + GroupNorm combination\n    },\n    \"dataset\": {\n        \"spatial_downsample\": [5,5],  # Reduce resolution\n        \"temporal_scale\": False,  # No temporal scaling (disabled as per request)\n        \"temporal_flip\":True,\n        \"temporal_shift\":False,\n        \"spatial_affine\":True,\n        \"frames_per_segment\": 50,  # Number of event frames\n        \"events_interpolation\": \"causal_linear\",  # Use all frames\n        \"time_window\": 10000\n    },\n    \"trainer\": {\n        \"batch_size\": 32,\n        \"learning_rate\": 0.001,\n        \"weight_decay\": 0.0002,\n        \"epochs\": 100,\n        \"device\": 0,\n        \"activity_regularization\" : 0\n    },\n    \"submission\": {\n        \"checkpoint_path\": \"weights/submission.ckpt\",\n        \"output_csv\": \"submission.csv\"\n    }\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:23:53.886678Z","iopub.execute_input":"2025-03-13T05:23:53.887024Z","iopub.status.idle":"2025-03-13T05:23:53.893221Z","shell.execute_reply.started":"2025-03-13T05:23:53.886995Z","shell.execute_reply":"2025-03-13T05:23:53.892238Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Eye Dataset\nimport ast\nimport math\nfrom pathlib import Path\n\nimport h5py\nimport numpy as np\nimport torch\nfrom natsort import natsorted\nfrom torch.nn import functional as F\nfrom torch.utils.data import Dataset\n\nrand_range = lambda amin, amax: amin + (amax - amin) * np.random.rand()\n\nval_files = [\"1_6\", \"2_4\", \"4_4\", \"6_2\", \"7_4\", \"9_1\", \"10_3\", \"11_2\", \"12_3\"]\n\ndef get_index(file_lens, index):\n    file_lens_cumsum = np.cumsum(np.array(file_lens))\n    file_id = np.searchsorted(file_lens_cumsum, index, side='right')\n    sample_id = index - file_lens_cumsum[file_id - 1] if file_id > 0 else index\n    return file_id, sample_id\n\n\ndef txt_to_npy(file_path):\n    data = []\n    with open(file_path, 'r') as file:\n        for line in file:\n            data.append(ast.literal_eval(line.strip()))\n    return np.array(data)\n\n\ndef h5_to_npy(file_path, name):\n    with h5py.File(file_path, 'r') as file:\n        npy_data = file[name][:]\n    return npy_data\n\n\ndef events_to_frames(events, size, num_frames, spatial_downsample, temporal_downsample, mode='bilinear'):\n    \"\"\"Perform bilinear interpolation directly on the events, \n    while converting them to frames and do spatial and temporal downsamplings \n    all at the same time.\n    \"\"\"\n    height, width = size    \n    p, x, y, t = events\n    events_frames = torch.zeros([num_frames, 2, height, width]).type_as(events)\n    \n    def bilinear_interp(x, scale, x_max):\n        if scale == 1:\n            return x, x, torch.ones_like(x), torch.zeros_like(x)\n        xd1 = x % scale / scale\n        xd = 1 - xd1\n        x = (x / scale).long().clamp(0, x_max)\n        x1 = (x + 1).clamp(0, x_max)\n        return x, x1, xd, xd1\n    \n    if mode == 'nearest':\n        p = p.round().long()\n        x = (x / spatial_downsample[0]).round().long().clamp(0, width - 1)\n        y = (y / spatial_downsample[1]).round().long().clamp(0, height - 1)\n        t = (t / temporal_downsample - 0.5).round().long().clamp(0, num_frames - 1)\n        events_frames.index_put_((t, p, y, x), torch.ones_like(p, dtype=torch.float32), accumulate=True)\n        return events_frames\n    \n    x, x1, xd, xd1 = bilinear_interp(x, spatial_downsample[0], width - 1)\n    y, y1, yd, yd1 = bilinear_interp(y, spatial_downsample[1], height - 1)\n    t, t1, td, td1 = bilinear_interp(t, temporal_downsample, num_frames - 1)\n    \n    # similar to bilinear, but temporally causal\n    if mode == 'causal_linear':\n        p = p.long().repeat(4)\n        \n        x = torch.cat([x.repeat(2), x1.repeat(2)])\n        y = torch.cat([y, y1]).repeat(2)\n        t = t.repeat(4)\n        \n        xd = torch.cat([xd.repeat(2), xd1.repeat(2)])\n        yd = torch.cat([yd, yd1]).repeat(2)\n        td = td1.repeat(4)  # causal\n        \n        events_frames.index_put_((t, p, y, x), xd * yd * td, accumulate=True)\n        return events_frames\n\n    # bilinear\n    p = p.long().repeat(8)\n\n    x = torch.cat([x.repeat(4), x1.repeat(4)])\n    y = torch.cat([y.repeat(2), y1.repeat(2)]).repeat(2)\n    t = torch.cat([t, t1]).repeat(4)\n\n    xd = torch.cat([xd.repeat(4), xd1.repeat(4)])\n    yd = torch.cat([yd.repeat(2), yd1.repeat(2)]).repeat(2)\n    td = torch.cat([td, td1]).repeat(4)\n\n    events_frames.index_put_((t, p, y, x), xd * yd * td, accumulate=True)\n    return events_frames\n\n# def events_to_frames(events, size, num_frames, spatial_downsample,temporal_downsample, mode='adaptive'):\n#     \"\"\"Adaptive binning based on event density.\"\"\"\n#     height, width = size    \n#     p, x, y, t = events\n\n#     # Compute event density per frame\n#     min_time, max_time = t.min(), t.max()\n#     frame_edges = torch.linspace(min_time, max_time, num_frames + 1)\n#     event_bins = torch.bucketize(t.to(frame_edges.device), frame_edges) - 1 # Assign events to bins\n\n#     events_frames = torch.zeros([num_frames, 2, height, width]).type_as(events)\n\n#     for i in range(num_frames):\n#         mask = (event_bins == i)\n#         if mask.sum() == 0:\n#             continue  # No events in this bin\n\n#         x_bin, y_bin, p_bin = x[mask], y[mask], p[mask]\n\n#         # Perform bilinear interpolation for spatial assignment\n#         x_low, x_high = x_bin.floor().long(), (x_bin + 1).floor().long()\n#         y_low, y_high = y_bin.floor().long(), (y_bin + 1).floor().long()\n\n#         x_low = x_low.clamp(0, width - 1)\n#         x_high = x_high.clamp(0, width - 1)\n#         y_low = y_low.clamp(0, height - 1)\n#         y_high = y_high.clamp(0, height - 1)\n\n#         # Assign events to bins with weights\n#         events_frames[i, p_bin.long(), y_low, x_low] += 1\n#         events_frames[i, p_bin.long(), y_low, x_high] += 1\n#         events_frames[i, p_bin.long(), y_high, x_low] += 1\n    #     events_frames[i, p_bin.long(), y_high, x_high] += 1\n\n    # return events_frames\n\n\n\nclass EventRandomAffine():\n    \"\"\"Perform random affine transformations on the events and labels\n    \"\"\"\n    def __init__(self, size, \n                 degrees=15, translate=(0.2, 0.2), scale=(0.8, 1.2), spatial_jitter=None, \n                 augment_flag=True):\n        self.degrees = degrees\n        self.translate = translate\n        self.scale = scale\n        self.spatial_jitter = spatial_jitter\n        self.augment_flag = augment_flag\n\n        self.height, self.width = size\n\n    def normalize(self, coords, backward=False):\n        if not backward:\n            coords[0] = coords[0] / self.width - 0.5\n            coords[1] = coords[1] / self.height - 0.5\n        else:\n            coords[0] = (coords[0] + 0.5) * self.width\n            coords[1] = (coords[1] + 0.5) * self.height\n\n        return coords\n\n    def __call__(self, events, labels):        \n        if self.augment_flag:\n            degrees = rand_range(-self.degrees, self.degrees) / 180 * math.pi\n            translate = [rand_range(-t, t) for t in self.translate]\n            scale = [rand_range(*self.scale) for _ in range(2)]\n\n            cos, sin = math.cos(degrees), math.sin(degrees)\n            R = torch.tensor([[cos, -sin, 0], [sin, cos, 0], [0, 0, 1]], dtype=torch.float).type_as(events)\n            S = torch.tensor([[scale[0], 0, 0], [0, scale[1], 0], [0, 0, 1]], dtype=torch.float).type_as(events)\n            T = torch.tensor([[1, 0, translate[0]], [0, 1, translate[1]], [0, 0, 1]], dtype=torch.float).type_as(events)\n\n            trans_matrix = T @ R @ S\n        \n        else:\n            trans_matrix = torch.eye(3).type_as(events)\n        \n        coords = F.pad(events[1:3], (0, 0, 0, 1), value=1)\n        coords = self.normalize(trans_matrix @ self.normalize(coords), True)\n        if self.spatial_jitter is not None:\n            coords += torch.randn_like(coords) * self.spatial_jitter\n        \n        events[1:3] = coords[:2]\n        val_inds = (coords[0] >= 0) & (coords[0] < self.width) & (coords[1] >= 0) & (coords[1] < self.height)\n        events = events[:, val_inds]\n        \n        labels = labels.T\n        centers = F.pad(labels[:2], (0, 0, 0, 1), value=1)\n        centers = trans_matrix @ self.normalize(centers)\n        centers = centers[:2] + 0.5\n        \n        closes = labels[-1]\n        \n        return events, centers, closes\n    \n    \nclass EyeTrackingDataset(Dataset):\n    def __init__(self, \n                 root_path, \n                 mode='train', \n                 device='cpu', \n                 time_window=10000, \n                 frames_per_segment=50, \n                 spatial_downsample=(5, 5), \n                 events_interpolation='bilinear', \n                 spatial_affine=True, \n                 temporal_flip=True, \n                 temporal_scale=True, \n                 temporal_shift=True,\n                 test_on_val=False):\n        self.mode = mode\n        self.time_window = time_window\n        self.frames_per_segment = frames_per_segment\n        self.time_window_per_segment = time_window * frames_per_segment\n        self.spatial_downsample = spatial_downsample\n        self.events_interpolation = events_interpolation\n        assert time_window == 10000\n        \n        self.temporal_flip = temporal_flip\n        self.temporal_scale = temporal_scale\n        self.temporal_shift = temporal_shift\n        \n        self.test_on_val = test_on_val\n        \n        root_path = Path(root_path)\n        if mode in ['train', 'val']:\n        #if mode == 'train':\n            base_path = root_path / 'train'\n        elif mode == 'test':\n            if test_on_val:\n                base_path = root_path / 'train'\n            else:\n                base_path = root_path / 'test'\n        else:\n            raise ValueError(\"Invalid mode. Most be train or test.\")\n        #data_dirs = natsorted(base_path.glob('*'))\n        \n        self.events, self.labels = [], []\n        self.num_frames_list, self.num_segments_list = [], []\n        \n        dir_paths = natsorted(base_path.glob('*'))\n        if mode == 'train':\n            dir_paths = [dir_path for dir_path in dir_paths if dir_path.name not in val_files]\n        elif mode == 'val' or (mode == 'test' and test_on_val):\n            dir_paths = [dir_path for dir_path in dir_paths if dir_path.name in val_files]\n\n\n        for dir_path in dir_paths:\n#        for dir_path in data_dirs:\n            assert dir_path.is_dir()\n            data_path = dir_path / f'{dir_path.name}.h5'\n            label_path = dir_path / 'label.txt' if (mode != 'test' or test_on_val) else dir_path / 'label_zeros.txt'\n            \n            event, label = h5_to_npy(data_path, 'events'), txt_to_npy(label_path)\n            \n            num_frames = label.shape[0]\n            self.num_frames_list.append(num_frames)\n            self.num_segments_list.append(num_frames // frames_per_segment)\n            \n            # truncating off trailing events with no labels\n            final_t = num_frames * time_window\n            final_ind = np.searchsorted(event['t'], final_t, 'left')\n            event = event[:final_ind]\n            \n            label = torch.tensor(label, dtype=torch.float, device=device)\n            event = np.stack([event['p'].astype('float32'), event['x'].astype('float32'), event['y'].astype('float32'), event['t'].astype('float32')], axis=0)\n            event = torch.tensor(event, dtype=torch.float, device=device)  # (4, N)\n            \n            self.events.append(event)\n            self.labels.append(label)\n        \n        self.total_segments = sum(self.num_segments_list)\n        \n        # spatial affine transformation\n        augment_flag = (mode == 'train') and spatial_affine\n        self.augment = EventRandomAffine((480, 640), augment_flag=augment_flag)\n            \n    def __len__(self):\n        if self.mode == 'test':\n            return len(self.events)\n        return self.total_segments\n    \n    def _process_data(self, event, label, index=None):\n        event, center, close = self.augment(event, label)\n        num_frames = self.frames_per_segment if self.mode != 'test' else self.num_frames_list[index]\n        \n        event = events_to_frames(event, \n                                 (480 // self.spatial_downsample[1], 640 // self.spatial_downsample[0]), \n                                 num_frames, self.spatial_downsample, self.time_window, \n                                 mode=self.events_interpolation)\n        \n        # time + polarity flip\n        if self.mode == 'train' and self.temporal_flip and np.random.rand() > 0.5:\n            event = event.flip(0).flip(1)  # (T, C, H, W)\n            center = center.flip(-1)\n            close = close.flip(-1)\n        \n        return event.moveaxis(0, 1), center, 1 - close\n    \n    def __getitem__(self, index):\n        if self.mode == 'test':\n            event, label = self.events[index], self.labels[index]\n            return self._process_data(event, label, index)\n        \n        file_id, segment_id = get_index(self.num_segments_list, index)\n        event, label = self.events[file_id], self.labels[file_id]\n        \n        start_t = segment_id * self.time_window * self.frames_per_segment\n        end_t = start_t + self.time_window * self.frames_per_segment\n        \n        # random temporal shift\n        max_offset = round(self.time_window_per_segment * 0.1)\n        if self.mode == 'train' and self.temporal_shift and start_t >= max_offset:\n            offset = np.random.rand() * max_offset\n            start_t -= offset\n            end_t -= offset\n        else:\n            offset = 0\n        \n        # random temporal scaling\n        num_frames = self.num_frames_list[file_id]\n        event = event.clone()\n        if self.mode == 'train' and self.temporal_scale and end_t < (num_frames * self.time_window * 0.8):\n            # scale_factor = float(rand_range(0.8, 1.2))\n            # event[-1] *= scale_factor\n            scale_factor = 1\n        else:\n            scale_factor = 1\n        \n        start_ind = torch.searchsorted(event[-1], start_t, side='left')\n        end_ind = torch.searchsorted(event[-1], end_t, side='left')\n        \n        event_segment = event[:, start_ind.item():end_ind.item()]\n        event_segment[-1] -= start_t\n        \n        start_label_id = segment_id * self.frames_per_segment\n        end_label_id = (segment_id + 1) * self.frames_per_segment\n        \n        # label interpolation\n        label_numpy = label.cpu().numpy()\n        num_frame = label_numpy.shape[0]\n        arange = np.arange(0, num_frame)\n        label_offset = offset / self.time_window\n        interp_range = np.linspace(\n            (start_label_id - label_offset) / scale_factor, \n            (end_label_id - label_offset - 1) / scale_factor, \n            self.frames_per_segment, \n        )\n        x_interp = np.interp(interp_range, arange, label_numpy[:, 0])\n        y_interp = np.interp(interp_range, arange, label_numpy[:, 1])\n        closeness = label_numpy[start_label_id:end_label_id, -1]\n        label_segment = torch.tensor(np.stack([x_interp, y_interp, closeness], axis=1)).type_as(label)\n        \n        return self._process_data(event_segment, label_segment)\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:23:57.368496Z","iopub.execute_input":"2025-03-13T05:23:57.368828Z","iopub.status.idle":"2025-03-13T05:24:01.161849Z","shell.execute_reply.started":"2025-03-13T05:23:57.368804Z","shell.execute_reply":"2025-03-13T05:24:01.160955Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Losses\nimport copy\nimport math\n\nimport torch\nfrom torch.nn import functional as F\n\n\nclass OutputHook(list):\n    \"\"\" Hook to capture module outputs.\n    \"\"\"\n    def __call__(self, module, input, output):\n        self.append(output)\n    \n    \nclass MacsEstimationHook:\n    def __init__(self, num_conv_layers):\n        self.num_conv_layers = num_conv_layers\n        \n        self.layer_id = 0\n        self._params_per_layer = torch.zeros(num_conv_layers)\n        self._macs_per_layer = torch.zeros(num_conv_layers)\n        self._macs_per_layer_with_sparsity = torch.zeros(num_conv_layers)\n        \n        self.nonzeros = torch.zeros(num_conv_layers)\n        self.totals = torch.zeros(num_conv_layers)\n    \n    def __call__(self, module, input, output):\n        output_size = math.prod((output.shape[1],) + output.shape[3:])\n        macs_weight = output_size * math.prod(module.weight.shape[1:])\n        macs_bias = output_size\n        \n        self._params_per_layer[self.layer_id] = module.weight.numel() + output.shape[1]\n        \n        macs = macs_weight + macs_bias\n        self._macs_per_layer[self.layer_id] = macs\n        \n        self.nonzeros[self.layer_id] += (input[0] != 0).sum().item()\n        self.totals[self.layer_id] += input[0].numel()\n        self._macs_per_layer_with_sparsity[self.layer_id] = macs * self.nonzeros[self.layer_id] / self.totals[self.layer_id]\n        \n        self.layer_id = (self.layer_id + 1) % self.num_conv_layers\n        \n    @property\n    def macs_per_layer(self):\n        return self._macs_per_layer.round().long()\n    \n    @property\n    def macs_per_layer_with_sparsity(self):\n        return self._macs_per_layer_with_sparsity.round().long()\n    \n    @property\n    def params_per_layer(self):\n        return self._params_per_layer.round().long()\n\n\nclass RegularizationLoss():\n    def __init__(self, reg_factor, model):\n\n        self.reg_factor = reg_factor # 1e-1 was awesome!!!\n        if reg_factor > 0:\n            # Hook for regularization of activations of ReLUs\n            self.output_hook = OutputHook()\n            for mm in model.modules():\n                if isinstance(mm, torch.nn.ReLU):\n                    mm.register_forward_hook(self.output_hook)\n\n    def __call__(self, ):\n        if self.reg_factor > 0:\n            l1_penalty = 0.\n            for output in self.output_hook:\n                l1_penalty += torch.norm(output, 1)/output.numel()\n            l1_penalty *= self.reg_factor\n            self.output_hook.clear()\n            return l1_penalty\n        else:\n            return 0.\n\n\ndef regression_loss(pred, center, openness):\n    x, y = center.moveaxis(1, 0)\n    \n    pred = torch.sigmoid(pred).clamp(1e-4, 1 - 1e-4)\n    center_loss = F.smooth_l1_loss(pred, center, beta=0.11, reduction='none').sum(1)  # (batch, frames)\n    valid_mask = openness.eq(1) & x.gt(0) & x.lt(1) & y.gt(0) & y.lt(1)\n    center_loss = torch.where(valid_mask, center_loss, 0).mean()\n    \n    return center_loss\n\n\ndef tracking_loss(pred, center, openness, gamma=2):\n    device = pred.device\n    batch_size, _, frames, height, width = pred.shape\n    \n    x, y = center.moveaxis(1, 0)\n    x_ind = (x * width).long().clamp(0, width - 1)  # (batch, frames)\n    y_ind = (y * height).long().clamp(0, height - 1)\n    x_mod = (x * width) % 1\n    y_mod = (y * height) % 1\n    center_mod = torch.stack([x_mod, y_mod], dim=1)  # (batch, 2, frames)\n    \n    pred = torch.sigmoid(pred).clamp(1e-4, 1 - 1e-4)\n    pred_pupil, pred_center_mod = pred[:, 0], pred[:, 1:]\n    \n    valid_mask = openness.eq(1) & x.gt(0) & x.lt(1) & y.gt(0) & y.lt(1)\n    pupil_mask = torch.zeros_like(pred_pupil).bool()  # (batch, frames, height, width)\n    \n    batch_range = torch.arange(batch_size, device=device).repeat_interleave(frames)\n    frames_range = torch.arange(frames, device=device).repeat(batch_size)\n    pupil_mask[batch_range, frames_range, y_ind.flatten(), x_ind.flatten()] = 1\n    \n    # (batch, frames, height, width)\n    center_loss = F.smooth_l1_loss(pred_center_mod, center_mod[..., None, None], beta=0.11, reduction='none').sum(1)\n    \n    focal_loss = torch.where(\n        pupil_mask, \n        -1 * (1 - pred_pupil).pow(gamma) * pred_pupil.log() + center_loss, \n        -1 * pred_pupil.pow(gamma) * (1 - pred_pupil).log(), \n    )  # (batch, frames, height, width)\n    \n    return focal_loss[valid_mask].sum() / valid_mask.sum()\n\n\nclass Losses():\n    \"\"\" \n    Gathers the different losses\n    \"\"\"\n    def __init__(self, detector_head, reg_factor, model):\n        self.prediction_loss = tracking_loss if detector_head else regression_loss\n        self.regularization_loss = RegularizationLoss(reg_factor, model)\n\n    def __call__(self, pred, center, openness):\n        loss = self.prediction_loss(pred, center, openness)\n        loss += self.regularization_loss()\n        return loss\n\n\ndef process_detector_prediction(pred):\n    device = pred.device\n\n    if len(pred.shape)==3: # basic head case\n        batch_size, _, frames = pred.shape\n        x = torch.sigmoid(pred[:,0,:])\n        y = torch.sigmoid(pred[:,1,:])\n        \n    else: # centernet head case\n        batch_size, _, frames, height, width = pred.shape\n        \n        pred_pupil, pred_x_mod, pred_y_mod = pred.moveaxis(1, 0)\n        pred_x_mod = torch.sigmoid(pred_x_mod)\n        pred_y_mod = torch.sigmoid(pred_y_mod)\n        \n        pupil_ind = pred_pupil.flatten(-2, -1).argmax(-1)  # (batch, frames)\n        pupil_ind_x = pupil_ind % width\n        pupil_ind_y = pupil_ind // width\n        \n        batch_range = torch.arange(batch_size, device=device).repeat_interleave(frames)\n        frames_range = torch.arange(frames, device=device).repeat(batch_size)\n        \n        pred_x_mod = pred_x_mod[batch_range, frames_range, pupil_ind_y.flatten(), pupil_ind_x.flatten()]\n        pred_y_mod = pred_y_mod[batch_range, frames_range, pupil_ind_y.flatten(), pupil_ind_x.flatten()]\n\n        x = (pupil_ind_x + pred_x_mod.view(batch_size, frames)) / width\n        y = (pupil_ind_y + pred_y_mod.view(batch_size, frames)) / height\n    \n    return torch.stack([x, y], dim=1)\n    \n\ndef p10_acc(pred, center, openness, detector_head=True, \n            height=60, width=80, tolerance=10):\n    pred = pred.detach().clone()\n    center = center.detach().clone()\n    \n    if detector_head:\n        pred = process_detector_prediction(pred)\n    else:\n        pred = torch.sigmoid(pred)\n    \n    pred[:, 0] *= width\n    pred[:, 1] *= height\n    center[:, 0] *= width\n    center[:, 1] *= height\n    \n    distances = torch.norm(center - pred, dim=1)\n    distances_noblinks = distances[openness == 1]\n\n    return (distances < tolerance).sum() / distances.numel(), (distances_noblinks < tolerance).sum() / distances_noblinks.numel(), distances.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:24:01.163109Z","iopub.execute_input":"2025-03-13T05:24:01.163575Z","iopub.status.idle":"2025-03-13T05:24:01.186274Z","shell.execute_reply.started":"2025-03-13T05:24:01.163545Z","shell.execute_reply":"2025-03-13T05:24:01.185142Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#generate val sub\nimport os\nfrom pathlib import Path\n\nfrom omegaconf import OmegaConf as OC\n\nimport numpy as np\nimport torch\nimport matplotlib.pyplot as plt\n\n# from eye_dataset import EyeTrackingDataset\n# from tenn_model import TennSt\n# from losses import process_detector_prediction, OutputHook, MacsEstimationHook\n\ntorch.set_grad_enabled(False)\n\ndef check_val_score(checkpoint_path, checkpoint_config, remove_blinks=False, test_on_val=True):\n    data_path = '/kaggle/input/event-based-eye-tracking-cvpr-2025/event_data/event_data'\n\n    if test_on_val:\n        # Val Data\n        data_files = [\"1_6\", \"2_4\", \"4_4\", \"6_2\", \"7_4\", \"9_1\", \"10_3\", \"11_2\", \"12_3\"]\n    else:\n        # Test data\n        data_files = [\"1_1\", \"2_2\", \"3_1\", \"4_2\", \"5_2\", \"6_4\", \"7_5\", \"8_2\", \"8_3\", \"10_2\", \"11_3\", \"12_4\"]\n\n    print(checkpoint_config)\n    print(checkpoint_path)\n\n    config = OC.load(checkpoint_config)\n    model = TennSt(**OC.to_container(config.model))\n    model.eval()\n    \n    if checkpoint_path is not None:\n        weights = torch.load(checkpoint_path, map_location='cpu')['state_dict']\n        print(list(weights.keys()))\n        # mystr = list(weights.keys())[0].split('backbone')[0]\n        # mysr = list(weights.keys())[0].split('gru')[0] # get the str before backbone\n        # weights = {k.partition(mystr)[2]: v for k, v in weights.items() if k.startswith((mystr,mysr))}\n        # model.load_state_dict(weights)\n        # Load checkpoint\n        # weights = torch.load(checkpoint_path, map_location='cpu')['state_dict']\n        \n        # Find the common prefix (before 'backbone' or 'gru')\n        prefix_backbone = list(weights.keys())[0].split('backbone')[0]\n        prefix_gru = list(weights.keys())[0].split('gru')[0]\n        \n        # Remove the prefixes and correctly load matching keys\n        filtered_weights = {}\n        for k, v in weights.items():\n            if k.startswith(prefix_backbone):\n                new_key = k[len(prefix_backbone):]  # Remove the prefix\n            elif k.startswith(prefix_gru):\n                new_key = k[len(prefix_gru):]  # Remove the prefix\n            else:\n                continue  # Skip unrelated keys\n            filtered_weights[new_key] = v\n        \n        # Load weights into model\n        model.load_state_dict(filtered_weights,strict=False)  # Allow non-matching keys\n\n\n    testset = EyeTrackingDataset(data_path, 'test', **OC.to_container(config.dataset), test_on_val=test_on_val)\n\n    collected_distances = np.zeros((0,))\n\n    # Setup per hooks to grab outputs from ReLU layers to measure sparsity\n    output_hook = OutputHook()\n    relu_counter = 0\n    for mm in model.modules():\n        if isinstance(mm, torch.nn.ReLU):\n            mm.register_forward_hook(output_hook)\n            relu_counter+=1\n\n    evdensity_per_layer = []\n    for ll in range(relu_counter):\n        evdensity_per_layer.append({\n            'evs_per_time': np.zeros((0, ))\n        })\n        \n    num_conv_layers = 0\n    for mm in model.modules():\n        if isinstance(mm, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)):\n            num_conv_layers += 1\n        \n    macs_hook = MacsEstimationHook(num_conv_layers)\n    for mm in model.modules():\n        if isinstance(mm, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)):\n            mm.register_forward_hook(macs_hook)\n    \n    results = []\n    datafile = 0\n    for (event, center, openness) in testset:\n        pred = model(event.unsqueeze(0))\n        pred = process_detector_prediction(pred).squeeze(0)\n\n        # Grab layer outputs for sparsity calculations\n        for ll, outs in enumerate(output_hook):\n            summation_axis = (0, 1)\n            for dim in range(3, outs.ndim):\n                summation_axis += (dim,)\n            \n            evs = torch.sum(outs > 0, axis=summation_axis).detach().numpy() / (np.prod(outs.shape) / outs.shape[2])\n            evdensity_per_layer[ll]['evs_per_time'] = np.concatenate((evdensity_per_layer[ll]['evs_per_time'], evs))\n        output_hook.clear()\n        \n        pred[0] *= 80\n        pred[1] *= 60\n        center[0] *= 80\n        center[1] *= 60\n        distances = torch.norm(center - pred, dim=0)\n        if remove_blinks:\n            distances = distances[openness == 1]\n\n        pred = pred.detach().numpy()\n        center = center.detach().numpy()\n        distances = distances.detach().numpy()\n\n        p10 = (distances < 10).sum() / distances.size\n        distances_5th = distances[::5]\n        p10_5th = (distances_5th < 10).sum() / distances_5th.size\n        distances.mean()\n        collected_distances = np.concatenate((collected_distances, distances_5th), axis=0)\n        results.append(\n            {\n                'datafile': data_files[datafile],\n                'pred': pred,\n                'center': center,\n                'distances': distances,\n                'p10_all': p10,\n                'p10_5th': p10_5th,\n                'openness': openness\n            }\n        )\n        datafile += 1\n\n    metrics = {}\n    p10_total = (collected_distances < 10).sum() / collected_distances.size\n    print('Overall p10 (100Hz): ' + str(p10_total))\n    euc_total = collected_distances.mean()\n    print('Overall Euc. Dist (100Hz): ' + str(euc_total))\n    collected_5th_distances = collected_distances[::5]\n    p10_5th_total = (collected_5th_distances < 10).sum() / collected_5th_distances.size\n    print('Overall p10  (20Hz): ' + str(p10_5th_total))\n    metrics['p10'] = p10_5th_total\n    euc_5th_total = collected_5th_distances.mean()\n    print('Overall Euc. Dist (20Hz): ' + str(euc_5th_total))\n    metrics['distance'] = euc_5th_total\n    \n    mean_event_densities = []\n    for ll in evdensity_per_layer:\n        mean_event_densities.append(np.mean(ll['evs_per_time']))\n        \n    metrics['mean_event_density'] = mean_event_densities\n    metrics['macs_per_layer'] = list(macs_hook.macs_per_layer.numpy())\n    \n    print(f\"Parameters per conv layer: {macs_hook.params_per_layer}\")\n    print(f\"Parameters of the network: {sum(macs_hook.params_per_layer)}\")\n    \n    print(f\"\\nOutput event density per ReLU layer: {[float(f'{density:.3f}') for density in mean_event_densities]}\")\n    print(f\"MACs per frame for each conv layer: {macs_hook.macs_per_layer}\")\n    print(f\"Total MACs per frame of the network: {sum(macs_hook.macs_per_layer)}\")\n    print(f\"MACs per frame for each conv layer (considering sparsity): {macs_hook.macs_per_layer_with_sparsity}\")\n    print(f\"Total MACs per frame of the network (considering sparsity): {sum(macs_hook.macs_per_layer_with_sparsity)}\")\n    \n    return results, metrics\n\n\ndef plot_results(grouped_results, test_on_val=True):\n    outdir = './val_results' if test_on_val else './test_results'\n    os.makedirs(outdir, exist_ok = True) \n    refres = grouped_results[0]\n    for eix in range(len(refres)):\n        fig, axs = plt.subplots(3, 1, figsize=(12, 12), constrained_layout=True, sharex=True)\n        for rix, results in enumerate(grouped_results):\n            expt = results[eix]\n            distances = expt['distances']\n            axs[0].plot(distances)\n            if np.any(distances>10):\n                misses = np.where(distances>10)\n                axs[0].plot(misses, np.ones_like(misses)*(11+rix), '.r')\n\n            axs[1].plot(expt['pred'][0], alpha=0.3)\n            axs[2].plot(expt['pred'][1], alpha=0.3)\n\n        blinks = np.where(expt['openness']==0)\n        axs[0].plot(blinks, np.ones_like(blinks)*(-1), '.k')\n        axs[0].set_ylabel('Distance')\n        axs[0].set_title('Validation File: ' + expt['datafile'])\n        axs[0].set_ylim([-1.4, 20])\n\n        axs[1].plot(expt['center'][0], 'xkcd:aqua', label='X', linewidth=3)\n        axs[1].plot(blinks, np.ones_like(blinks)*(5), '.k')\n        axs[1].set_ylabel('X Position')\n        axs[1].set_ylim([0, 80])\n\n        axs[2].plot(expt['center'][1], 'xkcd:salmon', label='Y', linewidth=3)\n        axs[2].plot(blinks, np.ones_like(blinks)*(5), '.k')\n        axs[2].set_ylabel('Y Position')\n        axs[2].set_xlabel('Timestep')\n        axs[2].set_ylim([0, 60])\n        fig.savefig(os.path.join(outdir, 'results_'+expt['datafile']+'.png'))\n\n\n# if __name__=='__main__':\n    \n#     os.environ['CUDA_VISIBLE_DEVICES'] = ''\n#     torch.set_grad_enabled(False)\n    \n#     config = '/kaggle/working/config.yaml'\n#     checkpoints = '/kaggle/working/lightning_logs/version_0/checkpoints/last.ckpt'\n\n#     # centernet version\n#     # checkpoints = ['/home/scrouzet/AIS2024_CVPR/train_tenn/outputs/2024-03-22/11-05-58/lightning_logs/version_0/checkpoints/last.ckpt']\n#     # config = ['/home/scrouzet/AIS2024_CVPR/train_tenn/outputs/2024-03-22/11-05-58/lightning_logs/version_0/config.yaml']\n    \n#     # basic version\n#     # checkpoints = ['/home/scrouzet/AIS2024_CVPR/train_tenn/outputs/2024-03-22/11-42-18/lightning_logs/version_0/checkpoints/last.ckpt']\n#     # config = ['/home/scrouzet/AIS2024_CVPR/train_tenn/outputs/2024-03-22/11-42-18/lightning_logs/version_0/config.yaml']\n    \n#     test_on_val = True\n#     grouped_results = []\n#     for k, checkpoint in enumerate(checkpoints):\n#         results, _ = check_val_score(checkpoint_path=checkpoint,\n#                                     checkpoint_config = config[k],\n#                                     remove_blinks=False,\n#                                     test_on_val=test_on_val)\n#         grouped_results.append(results)        \n\n#     plot_results(grouped_results,\n#                 test_on_val=test_on_val)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:24:01.187976Z","iopub.execute_input":"2025-03-13T05:24:01.188279Z","iopub.status.idle":"2025-03-13T05:24:01.449175Z","shell.execute_reply.started":"2025-03-13T05:24:01.188256Z","shell.execute_reply":"2025-03-13T05:24:01.448427Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# SpatioTemporal Model\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\n\nimport warnings\n\n# warnings.formatwarning = lambda message, category, filename, lineno, line=None: \\\n#     f'{category.__name__}: {message}\\n'\n\nclass CausalGroupNorm(nn.GroupNorm):\n    \"\"\"A GroupNorm that does not use temporal statistics, to ensure causality\n    \"\"\"\n    def __init__(self, num_groups, num_channels, **kwargs):\n        super().__init__(num_groups, num_channels, **kwargs)\n        \n    def forward(self, input):\n        x = input.moveaxis(1, 2)  # (B, T, C, H, W)\n        x_shape = x.shape\n        x = x.flatten(0, 1)  # (B * T, C, H, W)\n        x = super().forward(x).reshape(x_shape)\n        return x.moveaxis(1, 2)  # (B, C, T, H, W)\n\n\nact_layer = lambda: nn.ReLU()\nbn_block = lambda features: nn.Sequential(nn.BatchNorm3d(features), act_layer())\ngn_block = lambda features: nn.Sequential(CausalGroupNorm(4, features), act_layer())\npw_conv = lambda in_channels, out_channels: nn.Conv3d(in_channels, out_channels, 1, bias=False)\n\n\nclass SpatialBlock(nn.Module):\n    def __init__(self, \n                 in_channels, \n                 out_channels, \n                 depthwise=False, \n                 kernel_size=1,\n                 full_conv3d=False, \n                 norms='mixed'):\n        super().__init__()\n        kernel = (kernel_size,3,3)\n        self.kernel_size = kernel_size\n        self.full_conv3d = full_conv3d\n        self.norms = norms\n        self.streaming_mode = False\n        self.fifo = None  # for streaming inference\n\n        if self.norms=='all_gn':\n            norm_block = gn_block\n        else :\n            norm_block = bn_block\n\n        if depthwise:\n            self.block = nn.Sequential(\n                nn.Conv3d(in_channels, in_channels, kernel, (1, 2, 2), (0, 1, 1), groups=in_channels, bias=False), \n                norm_block(in_channels), \n                pw_conv(in_channels, out_channels), \n                norm_block(out_channels), \n            )\n            \n        else:\n            self.block = nn.Sequential(\n                nn.Conv3d(in_channels, out_channels, kernel, (1, 2, 2), (0, 1, 1), bias=False), \n                norm_block(out_channels), \n            )\n        \n    def streaming(self, enabled=True):\n        if enabled:\n            assert not self.training, \"Can only use streaming mode during evaluation.\"\n        self.streaming_mode = enabled\n        \n    def reset_memory(self):\n        self.fifo = None\n    \n    def forward(self, input):\n        if self.full_conv3d: \n            if self.streaming_mode:\n                return self._streaming_forward(input)\n            input = F.pad(input, (0, 0, 0, 0, self.kernel_size - 1, 0))\n            return self.block(input)\n        else:         \n            return self.block(input)\n            \n    def _streaming_forward(self, input):\n        if self.fifo is None:\n            self.fifo = torch.zeros(*input.shape[:2], self.kernel_size, *input.shape[3:]).type_as(input)\n        self.fifo = torch.cat([self.fifo[:, :, 1:], input], dim=2)\n        return self.block(self.fifo)\n\n\nclass TemporalBlock(nn.Module):\n    def __init__(self, \n                 in_channels, \n                 out_channels, \n                 kernel_size=3, \n                 depthwise=False,\n                 full_conv3d=False, \n                 norms='mixed'):\n        super().__init__()\n        assert out_channels % 4 == 0  # needed for group norm to work\n        self.kernel_size = kernel_size\n        self.depthwise = depthwise\n        self.norms = norms\n        kernel = (kernel_size,3,3) if full_conv3d else (kernel_size,1,1)\n        \n        self.streaming_mode = False\n        self.fifo = None  # for streaming inference\n        \n        if self.norms=='mixed':\n            norm1_block = bn_block\n            norm2_block = gn_block\n        elif self.norms=='all_bn':\n            norm1_block = bn_block\n            norm2_block = bn_block\n        elif self.norms=='all_gn':\n            norm1_block = gn_block\n            norm2_block = gn_block\n\n        if depthwise:\n            self.block = nn.Sequential(\n                nn.Conv3d(in_channels, in_channels, kernel, groups=in_channels, bias=False), \n                norm1_block(in_channels), \n                pw_conv(in_channels, out_channels), \n                norm2_block(out_channels), \n            )\n            \n        else:\n            self.block = nn.Sequential(\n                nn.Conv3d(in_channels, out_channels, kernel, bias=False), \n                norm2_block(out_channels), \n            )\n\n    def streaming(self, enabled=True):\n        if enabled:\n            assert not self.training, \"Can only use streaming mode during evaluation.\"\n        self.streaming_mode = enabled\n        \n    def reset_memory(self):\n        self.fifo = None\n    \n    def forward(self, input):\n        if self.streaming_mode:\n            return self._streaming_forward(input)\n                  \n        input = F.pad(input, (0, 0, 0, 0, self.kernel_size - 1, 0))\n        return self.block(input)\n    \n    def _streaming_forward(self, input):\n        if self.fifo is None:\n            self.fifo = torch.zeros(*input.shape[:2], self.kernel_size, *input.shape[3:]).type_as(input)\n        self.fifo = torch.cat([self.fifo[:, :, 1:], input], dim=2)\n        return self.block(self.fifo)\n        \n# from mamba_ssm import Mamba\n\n# class MambaLayer(nn.Module):\n#     \"\"\"Replaces GRU with Mamba-based temporal modeling.\"\"\"\n#     def __init__(self, input_dim, hidden_dim, seq_len):\n#         super().__init__()\n#         self.mamba = Mamba(\n#             d_model=hidden_dim,  \n#             d_state=64,  # State size for SSM\n#             d_conv=4,  \n#             expand=2\n#         )\n#         self.projection = nn.Linear(hidden_dim, input_dim)  # Adjust output\n\n#     def forward(self, x):\n#         # x: (batch, seq_len, channels, height, width)\n#         batch_size, seq_len, channels, height, width = x.shape\n#         x = x.mean((-2, -1))  # Reduce spatial dimensions -> (B, C, T)\n#         x = x.permute(0, 2, 1)  # Convert to (B, T, C)\n        \n#         # Check if the embedding dimension matches Mamba's input\n#         # if x.shape[-1] != self.mamba.input_dim:\n#         # x = self.projection(x)  # Apply linear projection if needed\n        \n#         x = self.mamba(x)  # Apply Mamba SSM\n#         x = x.permute(0, 2, 1)  # Convert back to (B, C, T)\n\n#         return x.view(batch_size, seq_len, channels)\n\n\nclass TennSt(nn.Module):\n    def __init__(\n        self, \n        channels, \n        t_kernel_size, \n        n_depthwise_layers, \n        detector_head, \n        detector_depthwise, \n        full_conv3d=False,\n        norms='mixed',\n    ):\n        super().__init__()\n        self.detector = detector_head\n        \n        depthwises = [False] * (10 - n_depthwise_layers) + [True] * n_depthwise_layers\n        temporals = [True, False] * 5\n\n        # self.mamba_layer = MambaLayer(input_dim=channels[-1], hidden_dim=256,seq_len=channels[-1])  \n        self.lstm = nn.LSTM(input_size=channels[-1], hidden_size=256, num_layers=2, batch_first=True)\n        \n        self.backbone = nn.Sequential()\n        for i in range(len(depthwises)):\n            in_channels, out_channels = channels[i], channels[i+1]\n            depthwise = depthwises[i]\n            temporal = temporals[i]\n            \n            if temporal:\n                self.backbone.append(TemporalBlock(in_channels, out_channels, \n                                                   kernel_size=t_kernel_size, depthwise=depthwise,\n                                                   full_conv3d=full_conv3d, norms=norms))\n            else:\n                self.backbone.append(SpatialBlock(in_channels, out_channels, depthwise=depthwise,\n                                                  full_conv3d=full_conv3d,\n                                                  kernel_size=t_kernel_size if full_conv3d else 1,\n                                                  norms=norms))\n        \n        if detector_head:\n            self.head = nn.Sequential(\n                TemporalBlock(channels[-1], channels[-1], t_kernel_size, depthwise=detector_depthwise), \n                nn.Conv3d(channels[-1], channels[-1], (1, 3, 3), (1, 1, 1), (0, 1, 1)), \n                act_layer(), \n                nn.Conv3d(channels[-1], 3, 1), \n            )\n        else:\n            self.head = nn.Sequential(\n                nn.Conv1d(channels[-1], channels[-1], 1), \n                act_layer(), \n                nn.Conv1d(channels[-1], 2, 1), \n            )\n    \n    def streaming(self, enabled=True):\n        if enabled:\n            warnings.warn(\"You have enabled the streaming mode of the network. It is expected, but not checked, that the input will be of shape (batch, 1, H, W).\")\n        for name, module in self.named_modules():\n            if name and hasattr(module, 'streaming'):\n                module.streaming(enabled)\n                \n    def reset_memory(self):\n        for name, module in self.named_modules():\n            if name and hasattr(module, 'reset_memory'):\n                module.reset_memory()\n         \n    def forward(self, input: torch.Tensor) -> torch.Tensor:\n        # if self.detector:\n        #     return self.head((self.backbone(input)))\n        # else:\n        #     return self.head(self.backbone(input).mean((-2, -1)))\n\n        x = self.backbone(input)  # (B, C, T, H, W)\n        x = x.mean((-2, -1))  # Reduce spatial dimensions -> (B, C, T)\n        \n        # Prepare for GRU: Convert to (B, T, C)\n        x = x.permute(0, 2, 1)  # (B, T, C)\n        # x = self.mamba_layer(x)\n        x, _ = self.lstm(x)  # Apply GRU\n        x = x.permute(0, 2, 1)  # Convert back to (B, C, T)\n        \n        # Expand spatial dimensions back before passing to Conv3D\n        x = x.unsqueeze(-1).unsqueeze(-1)  # (B, C, T, 1, 1)\n        \n        return self.head(x) if self.detector else self.head(x.mean(-2,-1))\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:26:10.017084Z","iopub.execute_input":"2025-03-13T05:26:10.017602Z","iopub.status.idle":"2025-03-13T05:26:10.044214Z","shell.execute_reply.started":"2025-03-13T05:26:10.017554Z","shell.execute_reply":"2025-03-13T05:26:10.043188Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:26:10.392047Z","iopub.execute_input":"2025-03-13T05:26:10.392376Z","iopub.status.idle":"2025-03-13T05:26:10.396684Z","shell.execute_reply.started":"2025-03-13T05:26:10.392349Z","shell.execute_reply":"2025-03-13T05:26:10.395615Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport csv\nfrom datetime import datetime\nfrom functools import partial\nfrom pathlib import Path\n\nimport hydra\nimport torch\nimport torch._dynamo\ntorch._dynamo.config.suppress_errors = True\n\nfrom lightning import LightningModule, Trainer\nfrom lightning.pytorch.callbacks import ModelCheckpoint\nfrom omegaconf import OmegaConf as OC\nfrom pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\n\nconfig = OC.create(config)\n\n# import losses\n# from eye_dataset import EyeTrackingDataset\n# from tenn_model import TennSt\n# from generate_val_results import check_val_score\n\nfilename = 'results_ablation.csv'\n\ndef flatten_2levels_dict(d, parent_key='', sep='_'):\n    \"\"\"\n    Flatten a nested dictionary.\n    \"\"\"\n    items = {}\n    for pk, pv in d.items():\n        parent_key = pk + sep\n        for ck, cv in pv.items():\n            items[pk+'__'+ck] = cv\n    return items\n\n\nclass CustomModule(LightningModule):\n    def __init__(self, data_path, config):\n        super().__init__()\n        self.data_path = data_path\n        self.config = config\n        \n        self.batch_size = config[\"trainer\"]['batch_size']\n        epochs = config['trainer']['epochs']\n        detector_head = config['model']['detector_head']\n        activity_regularization = config['trainer']['activity_regularization']\n        \n        self.model = TennSt(**OC.to_container(config.model))\n        \n        self.trainset = EyeTrackingDataset(data_path, 'train', config['trainer']['device'], **OC.to_container(config['dataset']))\n        self.valset = EyeTrackingDataset(data_path, 'val', config['trainer']['device'], **OC.to_container(config['dataset'])) # frames_per_segment=127, device=device)\n        \n        num_steps_per_epoch = len(self.trainset) // self.batch_size\n        self.total_train_steps = epochs * num_steps_per_epoch\n        \n        self.loss_fn = Losses(detector_head, activity_regularization, self.model)\n        self.metric_fn = partial(p10_acc, detector_head=detector_head)\n            \n    def forward(self, input):\n        return self.model(input)\n    \n    def _log(self, name, metric):\n        self.log(name, metric, \n                 on_step=False, on_epoch=True, prog_bar=True)\n    \n    def on_train_start(self):\n        log_dir = '/kaggle/working/'\n        OC.save(self.config, log_dir + 'config.yaml')\n\n    # Save final weights\n\n    #def on_train_end(self):\n        # torch.set_grad_enabled(False)\n        \n        # log_dir = Path(self.trainer.logger.log_dir)\n\n        # testset = EyeTrackingDataset(self.data_path, 'test', self.config.trainer.device, \n        #                              **OC.to_container(self.config.dataset))\n\n        # predictions = []\n        # for (event, _, _) in testset:\n        #     pred = self(event[None, ...])\n        #     if self.config.model.detector_head:\n        #         pred = losses.process_detector_prediction(pred)\n        #     else:\n        #         pred = torch.sigmoid(pred)\n        #     predictions.append(pred.detach().squeeze(0)[..., ::5].cpu().numpy())\n            \n        # predictions = np.concatenate(predictions, axis=-1).T  # (frames, 2)\n        # predictions[:, 0] *= 80\n        # predictions[:, 1] *= 60\n        # predictions = np.concatenate([np.arange(len(predictions))[:, None], predictions], axis=1)\n\n        # df = pd.DataFrame(predictions, columns=['row_id', 'x', 'y'])\n        # df.to_csv(log_dir / 'submission.csv', index=False)\n    \n    \n    def training_step(self, batch, batch_idx):\n        event, center, openness = batch\n        pred = self(event)        \n        loss = self.loss_fn(pred, center, openness)\n        metric, metric_noblinks, distance = self.metric_fn(pred, center, openness)\n        self._log('train_loss', loss)\n        self._log('train_metric', metric)\n        self._log('train_metric_noblinks', metric_noblinks)\n        self._log('train_distance', distance)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        event, center, openness = batch\n        pred = self(event)        \n        loss = self.loss_fn(pred, center, openness)\n        metric, metric_noblinks, distance = self.metric_fn(pred, center, openness)\n        self._log('val_loss', loss)\n        self._log('val_metric', metric)\n        self._log('val_metric_noblinks', metric_noblinks)\n        self._log('val_distance', distance)\n            \n    def configure_optimizers(self):\n        optimizer = AdamW(self.model.parameters(), lr=0.002, weight_decay=0.001)\n        scheduler = LinearWarmupCosineAnnealingLR(optimizer, round(self.total_train_steps * 0.025), self.total_train_steps, eta_min=1e-5)\n        \n        scheduler = {'scheduler': scheduler, \n                     'interval': 'step', \n                     'frequency': 1}\n        return [optimizer], [scheduler]\n            \n    def train_dataloader(self):\n        return DataLoader(self.trainset, shuffle=True, drop_last=True, batch_size=self.batch_size)\n\n    def val_dataloader(self):\n        return DataLoader(self.valset, shuffle=False, drop_last=False, batch_size=self.batch_size)\n    \n\n\ndef main():\n    data_path = '/kaggle/input/event-based-eye-tracking-cvpr-2025/event_data/event_data'\n    module = CustomModule(data_path, config)\n\n    checkpoint_callback = ModelCheckpoint(\n        monitor='val_metric', \n        mode='max', \n        save_last=True, \n        every_n_epochs=1, \n        filename='{epoch}-{val_metric:.2f}', \n    )\n\n\n\n    trainer = Trainer(\n        max_epochs=config['trainer']['epochs'], \n        gradient_clip_val=3.0, \n        accelerator='gpu', \n        devices=[config['trainer']['device']], \n        benchmark=True, \n        log_every_n_steps=5, \n        callbacks=[checkpoint_callback]\n    )\n\n    timestamp_start = datetime.now()\n    trainer.fit(module)\n    timestamp_end = datetime.now()\n\n    log_dir = Path(trainer.log_dir)\n    print(log_dir)\n    _, resmetrics = check_val_score(log_dir / \"checkpoints\" / \"last.ckpt\",\n                              \"/kaggle/working/config.yaml\",\n                              remove_blinks=False, test_on_val=True)\n    resdict = flatten_2levels_dict(config)\n    resdict['time_start'] = timestamp_start.isoformat()\n    resdict['time_end'] = timestamp_end.isoformat()\n    resdict['training_duration'] = int((timestamp_end-timestamp_start).total_seconds()/60) # training duration in minutes\n    resdict['params'] = sum(p.numel() for p in module.parameters())\n    resdict['params_trainable'] = sum(p.numel() for p in module.parameters() if p.requires_grad)\n\n    resdict['logdir'] = trainer.log_dir\n    resdict['val_p10'] = resmetrics['p10']\n    resdict['val_distance'] = resmetrics['distance']\n    resdict['event_density'] = resmetrics['mean_event_density']\n    resdict['macs'] = resmetrics['macs_per_layer']\n\n    # # Write the header and results to the CSV file\n    # with open(os.path.dirname('/kaggle/working/'), 'a', newline='') as csvfile:\n    #     writer = csv.DictWriter(csvfile, fieldnames=resdict.keys())\n    #     if csvfile.tell() == 0:\n    #         writer.writeheader()  # Write header only if file is empty\n    #     row = {**resdict}\n    #     writer.writerow(row)\n    # print(\"\\n\\n\")\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:26:10.933730Z","iopub.execute_input":"2025-03-13T05:26:10.934096Z","execution_failed":"2025-03-13T05:33:14.956Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom pathlib import Path\nimport time\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom omegaconf import OmegaConf as OC\n\n# from eye_dataset import EyeTrackingDataset\n# from tenn_model import TennSt\n# from losses import process_detector_prediction\n\n# NOTE: this submission script runs the network in streaming mode, and does not use the GPU\nos.environ['CUDA_VISIBLE_DEVICES'] = ''\n# torch.set_num_threads(1)\ntorch.set_grad_enabled(False)\n\n\ndef streaming_inference(model, frames):\n    model.eval()\n    model.streaming()\n    model.reset_memory()\n    \n    predictions = []\n    with torch.inference_mode():\n        for frame_id in range(frames.shape[2]):  # stream the frames to the model\n            prediction = model(frames[:, :, [frame_id]])\n            predictions.append(prediction)\n                \n    predictions = torch.cat(predictions, dim=2)\n    return predictions\n\n\nconfig_path = '/kaggle/working/config.yaml'\ncheckpoint_path = '/kaggle/working/lightning_logs/version_0/checkpoints/last.ckpt'\n\n#config_path = '/home/scrouzet/AIS2024_CVPR/train_tenn/outputs/2024-03-22/06-03-29/lightning_logs/version_0/config.yaml'\n#checkpoint_path = '/home/scrouzet/AIS2024_CVPR/train_tenn/outputs/2024-03-22/06-03-29/lightning_logs/version_0/checkpoints/last.ckpt'\n\nconfig = OC.load(config_path)\ndata_path = '/kaggle/input/event-based-eye-tracking-cvpr-2025/event_data/event_data'\n\nweights = torch.load(checkpoint_path, map_location='cpu')['state_dict']\nmystr = list(weights.keys())[0].split('backbone')[0] # get the str before backbone\nweights = {k.partition(mystr)[2]: v for k, v in weights.items() if k.startswith(mystr)}\n\nmodel = TennSt(**OC.to_container(config.model))\nmodel.eval()\nmodel.load_state_dict(weights)\n\ntestset = EyeTrackingDataset(data_path, 'test', **OC.to_container(config.dataset))\nevent_frames_list = [event_frames for (event_frames, _, _) in testset]\n\npredictions = []\nfor event_frames in event_frames_list:\n    pred = streaming_inference(model, event_frames[None, :])\n    pred = process_detector_prediction(pred)\n    predictions.append(pred.squeeze(0))\n    \npredictions = torch.cat(predictions, dim=-1)\n\npredictions[0] *= 80\npredictions[1] *= 60\n\npredictions_numpy = predictions.detach().numpy().T\npredictions_numpy = np.concatenate([np.arange(len(predictions_numpy))[:, None], predictions_numpy], axis=1)\n\ndf = pd.DataFrame(predictions_numpy, columns=['row_id', 'x', 'y'])\ndf.to_csv('/kaggle/working/submission.csv', index=False)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T05:16:27.945382Z","iopub.status.idle":"2025-03-13T05:16:27.945645Z","shell.execute_reply":"2025-03-13T05:16:27.945538Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}