Upload split_model.py with huggingface_hub

split_model.py

@@ -0,0 +1,651 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, Subset
+import torchvision.transforms as transforms
+from PIL import Image
+import os
+import numpy as np
+from bs4 import BeautifulSoup
+import argparse
+import logging
+from torch.utils.tensorboard import SummaryWriter
+from datetime import datetime
+import json
+from PIL import Image, ImageDraw
+import matplotlib.pyplot as plt
+
+
+def get_ground_truth(image, cells, otsl, split_width=5):
+
+    """
+    parse OTSL to derive row/column split positions.
+    this is the groundtruth for split model training.
+
+    Args:
+        image: PIL Image
+        html_tags: not used, kept for compatibility
+        cells: nested list - cells[0] contains actual cell data
+        otsl: OTSL token sequence
+        split_width: width of split regions in pixels (default: 5)
+    """
+    orig_width, orig_height = image.size
+    target_size = 960
+    
+    # cells is nested - extract actual list
+    cells_flat = cells[0]
+    
+    # parse OTSL to build 2D grid
+    grid = []
+    current_row = []
+    cell_idx = 0  # only increments for fcel ,ecel tokens
+    
+    for token in otsl:
+        if token == 'nl':
+            if current_row:
+                grid.append(current_row)
+                current_row = []
+        elif token == 'fcel' or token=='ecel':
+            current_row.append({'type': token, 'cell_idx': cell_idx})
+            cell_idx += 1
+        elif token in ['lcel', 'ucel', 'xcel']:
+            # merge/empty tokens don't consume bboxes
+            current_row.append({'type': token, 'cell_idx': None})
+    
+    if current_row:
+        grid.append(current_row)
+    
+    # derive row splits - max y2 for each row
+    row_splits = []
+    for row in grid:
+        row_cell_indices = [item['cell_idx'] for item in row if item['cell_idx'] is not None]
+        if row_cell_indices:
+            max_y = max(cells_flat[i]['bbox'][3] for i in row_cell_indices)
+            row_splits.append(max_y)
+    
+    # derive column splits - max x2 for each column
+    num_cols = len(grid[0]) if grid else 0
+    col_splits = []
+    for col_idx in range(num_cols):
+        col_max_x = []
+        for row in grid:
+            if col_idx < len(row) and row[col_idx]['cell_idx'] is not None:
+                next_is_lcel = (col_idx + 1 < len(row) and row[col_idx + 1]['type'] == 'lcel')
+                if not next_is_lcel:
+                    cell_id = row[col_idx]['cell_idx']
+                    col_max_x.append(cells_flat[cell_id]['bbox'][2])
+        if col_max_x:
+            col_splits.append(max(col_max_x))
+
+    # # DEBUG: print what we found
+    # print(f"\nDEBUG get_ground_truth:")
+    # print(f"  Found {len(row_splits)} row splits: {row_splits}")
+    # print(f"  Found {len(col_splits)} col splits: {col_splits}")
+    
+    # # scale to target size
+    # y_scaled = [(y * target_size / orig_height) for y in row_splits]
+    # x_scaled = [(x * target_size / orig_width) for x in col_splits]
+    
+    # print(f"  Scaled row splits: {[int(y) for y in y_scaled]}")
+    # print(f"  Scaled col splits: {[int(x) for x in x_scaled]}")
+
+    
+    row_splits = row_splits[:-1]
+    col_splits = col_splits[:-1]
+
+    # scale to target size
+    y_scaled = [(y * target_size / orig_height) for y in row_splits]
+    x_scaled = [(x * target_size / orig_width) for x in col_splits]
+    
+    # init ground truth arrays
+    horizontal_gt = [0] * target_size
+    vertical_gt = [0] * target_size
+
+    all_x1 = [c['bbox'][0] for c in cells_flat]
+    all_y1 = [c['bbox'][1] for c in cells_flat]
+    all_x2 = [c['bbox'][2] for c in cells_flat]
+    all_y2 = [c['bbox'][3] for c in cells_flat]
+    table_bbox = [min(all_x1), min(all_y1), max(all_x2), max(all_y2)]
+    table_y1 = int(round(table_bbox[1] * target_size / orig_height))
+    table_y2 = int(round(table_bbox[3] * target_size / orig_height))
+    table_x1 = int(round(table_bbox[0] * target_size / orig_width))
+    table_x2 = int(round(table_bbox[2] * target_size / orig_width))
+
+
+    # Mark table bbox boundaries (5 pixels wide)
+    # Top boundary
+    for offset in range(split_width):
+        pos = table_y1 + offset
+        if 0 <= pos < target_size:
+            horizontal_gt[pos] = 1
+
+    # Bottom boundary
+    for offset in range(split_width):
+        pos = table_y2 - offset
+        if 0 <= pos < target_size:
+            horizontal_gt[pos] = 1
+
+    # Left boundary
+    for offset in range(split_width):
+        pos = table_x1 + offset
+        if 0 <= pos < target_size:
+            vertical_gt[pos] = 1
+
+    # Right boundary
+    for offset in range(split_width):
+        pos = table_x2 - offset
+        if 0 <= pos < target_size:
+            vertical_gt[pos] = 1
+
+    # mark split regions (configurable pixel width)
+    for y in y_scaled:
+        y_int = int(round(y))
+        if 0 <= y_int < target_size:
+            for offset in range(split_width):
+                pos = y_int + offset
+                if 0 <= pos < target_size:
+                    horizontal_gt[pos] = 1
+
+    for x in x_scaled:
+        x_int = int(round(x))
+        if 0 <= x_int < target_size:
+            for offset in range(split_width):
+                pos = x_int + offset
+                if 0 <= pos < target_size:
+                    vertical_gt[pos] = 1
+    
+    return horizontal_gt, vertical_gt
+
+
+def get_ground_truth_auto_gap(image, cells, otsl):
+    """
+    Parse OTSL to derive row/column split positions with DYNAMIC gap widths.
+    This creates ground truth for the split model training.
+
+    Args:
+        image: PIL Image
+        cells: nested list - cells[0] contains actual cell data
+        otsl: OTSL token sequence
+    """
+    orig_width, orig_height = image.size
+    target_size = 960
+    
+    # cells is nested - extract actual list
+    cells_flat = cells[0]
+    
+    # Parse OTSL to build 2D grid
+    grid = []
+    current_row = []
+    cell_idx = 0  # only increments for fcel, ecel tokens
+    
+    for token in otsl:
+        if token == 'nl':
+            if current_row:
+                grid.append(current_row)
+                current_row = []
+        elif token == 'fcel' or token == 'ecel':
+            current_row.append({'type': token, 'cell_idx': cell_idx})
+            cell_idx += 1
+        elif token in ['lcel', 'ucel', 'xcel']:
+            # merge/empty tokens don't consume bboxes
+            current_row.append({'type': token, 'cell_idx': None})
+    
+    if current_row:
+        grid.append(current_row)
+    
+    # Get row boundaries (min y1 and max y2 for each row)
+    row_boundaries = []
+    for row in grid:
+        row_cell_indices = [item['cell_idx'] for item in row if item['cell_idx'] is not None]
+        if row_cell_indices:
+            min_y1 = min(cells_flat[i]['bbox'][1] for i in row_cell_indices)
+            max_y2 = max(cells_flat[i]['bbox'][3] for i in row_cell_indices)
+            row_boundaries.append({'min_y': min_y1, 'max_y': max_y2})
+    
+    # Get column boundaries (min x1 and max x2 for each column)
+    num_cols = len(grid[0]) if grid else 0
+    col_boundaries = []
+    for col_idx in range(num_cols):
+        col_cells = []
+        for row in grid:
+            if col_idx < len(row) and row[col_idx]['cell_idx'] is not None:
+                # Check if next cell is lcel (merged left)
+                next_is_lcel = (col_idx + 1 < len(row) and row[col_idx + 1]['type'] == 'lcel')
+                if not next_is_lcel:
+                    cell_id = row[col_idx]['cell_idx']
+                    col_cells.append(cell_id)
+        if col_cells:
+            min_x1 = min(cells_flat[i]['bbox'][0] for i in col_cells)
+            max_x2 = max(cells_flat[i]['bbox'][2] for i in col_cells)
+            col_boundaries.append({'min_x': min_x1, 'max_x': max_x2})
+    
+    # Calculate table bbox
+    all_x1 = [c['bbox'][0] for c in cells_flat]
+    all_y1 = [c['bbox'][1] for c in cells_flat]
+    all_x2 = [c['bbox'][2] for c in cells_flat]
+    all_y2 = [c['bbox'][3] for c in cells_flat]
+    table_bbox = [min(all_x1), min(all_y1), max(all_x2), max(all_y2)]
+    
+    # Init ground truth arrays
+    horizontal_gt = [0] * target_size
+    vertical_gt = [0] * target_size
+    
+    # Helper function to scale and mark range
+    def mark_range(gt_array, start, end, orig_dim):
+        """Mark all pixels from start to end (scaled to target_size)"""
+        start_scaled = int(round(start * target_size / orig_dim))
+        end_scaled = int(round(end * target_size / orig_dim))
+        for pos in range(start_scaled, min(end_scaled + 1, target_size)):
+            if 0 <= pos < target_size:
+                gt_array[pos] = 1
+    
+    # Mark HORIZONTAL gaps (between rows)
+    # 1. Gap from image top to first row top
+    if row_boundaries:
+        mark_range(horizontal_gt, 0, row_boundaries[0]['min_y'], orig_height)
+    
+    # 2. Gaps between consecutive rows
+    for i in range(len(row_boundaries) - 1):
+        gap_start = row_boundaries[i]['max_y']
+        gap_end = row_boundaries[i + 1]['min_y']
+        if gap_end > gap_start:  # Only mark if there's actual gap
+            mark_range(horizontal_gt, gap_start, gap_end, orig_height)
+    
+    # 3. Gap from last row bottom to image bottom
+    if row_boundaries:
+        mark_range(horizontal_gt, row_boundaries[-1]['max_y'], orig_height, orig_height)
+    
+    # Mark VERTICAL gaps (between columns)
+    # 1. Gap from image left to first column left
+    if col_boundaries:
+        mark_range(vertical_gt, 0, col_boundaries[0]['min_x'], orig_width)
+    
+    # 2. Gaps between consecutive columns
+    for i in range(len(col_boundaries) - 1):
+        gap_start = col_boundaries[i]['max_x']
+        gap_end = col_boundaries[i + 1]['min_x']
+        if gap_end > gap_start:  # Only mark if there's actual gap
+            mark_range(vertical_gt, gap_start, gap_end, orig_width)
+    
+    # 3. Gap from last column right to image right
+    if col_boundaries:
+        mark_range(vertical_gt, col_boundaries[-1]['max_x'], orig_width, orig_width)
+    
+    return horizontal_gt, vertical_gt
+
+
+def get_ground_truth_auto_gap_expand_min5pix_overlap_cells(image, cells, otsl, split_width=5):
+    """
+    Parse OTSL to derive row/column split positions with DYNAMIC gap widths.
+    This creates ground truth for the split model training.
+
+    Args:
+        image: PIL Image
+        cells: nested list - cells[0] contains actual cell data
+        otsl: OTSL token sequence
+        split_width: width of split when there's no gap (default: 5)
+    """
+    orig_width, orig_height = image.size
+    target_size = 960
+    
+    # cells is nested - extract actual list
+    cells_flat = cells[0]
+    
+    # Parse OTSL to build 2D grid
+    grid = []
+    current_row = []
+    cell_idx = 0  # only increments for fcel, ecel tokens
+    
+    for token in otsl:
+        if token == 'nl':
+            if current_row:
+                grid.append(current_row)
+                current_row = []
+        elif token in ['fcel', 'ecel']:  # FIXED: was == ['fcel','ecel']
+            current_row.append({'type': token, 'cell_idx': cell_idx})
+            cell_idx += 1
+        elif token in ['lcel', 'ucel', 'xcel']:
+            # merge/empty tokens don't consume bboxes
+            current_row.append({'type': token, 'cell_idx': None})
+    
+    if current_row:
+        grid.append(current_row)
+    
+    # Get row boundaries (min y1 and max y2 for each row)
+    row_boundaries = []
+    for row in grid:
+        row_cell_indices = [item['cell_idx'] for item in row if item['cell_idx'] is not None]
+        if row_cell_indices:
+            min_y1 = min(cells_flat[i]['bbox'][1] for i in row_cell_indices)
+            max_y2 = max(cells_flat[i]['bbox'][3] for i in row_cell_indices)
+            row_boundaries.append({'min_y': min_y1, 'max_y': max_y2, 'row_cells': row_cell_indices})
+    
+    # Get column boundaries (min x1 and max x2 for each column)
+    num_cols = len(grid[0]) if grid else 0
+    col_boundaries = []
+    for col_idx in range(num_cols):
+        col_cells = []
+        for row in grid:
+            if col_idx < len(row) and row[col_idx]['cell_idx'] is not None:
+                # Check if next cell is lcel (merged left)
+                next_is_lcel = (col_idx + 1 < len(row) and row[col_idx + 1]['type'] == 'lcel')
+                if not next_is_lcel:
+                    cell_id = row[col_idx]['cell_idx']
+                    col_cells.append(cell_id)
+        if col_cells:
+            min_x1 = min(cells_flat[i]['bbox'][0] for i in col_cells)
+            max_x2 = max(cells_flat[i]['bbox'][2] for i in col_cells)
+            col_boundaries.append({'min_x': min_x1, 'max_x': max_x2, 'col_cells': col_cells})
+    
+    # Calculate table bbox
+    all_x1 = [c['bbox'][0] for c in cells_flat]
+    all_y1 = [c['bbox'][1] for c in cells_flat]
+    all_x2 = [c['bbox'][2] for c in cells_flat]
+    all_y2 = [c['bbox'][3] for c in cells_flat]
+    table_bbox = [min(all_x1), min(all_y1), max(all_x2), max(all_y2)]
+    
+    # Init ground truth arrays
+    horizontal_gt = [0] * target_size
+    vertical_gt = [0] * target_size
+    
+    # Helper function to scale and mark range
+    def mark_range(gt_array, start, end, orig_dim):
+        """Mark all pixels from start to end (scaled to target_size)"""
+        start_scaled = int(round(start * target_size / orig_dim))
+        end_scaled = int(round(end * target_size / orig_dim))
+        for pos in range(start_scaled, min(end_scaled + 1, target_size)):
+            if 0 <= pos < target_size:
+                gt_array[pos] = 1
+    
+    # Mark HORIZONTAL gaps (between rows)
+    # 1. Gap from image top to first row top
+    if row_boundaries:
+        mark_range(horizontal_gt, 0, row_boundaries[0]['min_y'], orig_height)
+    
+    # 2. Gaps between consecutive rows
+    for i in range(len(row_boundaries) - 1):
+        gap_start = row_boundaries[i]['max_y']
+        gap_end = row_boundaries[i + 1]['min_y']
+        if gap_end > gap_start:  # Only mark if there's actual gap
+            mark_range(horizontal_gt, gap_start, gap_end, orig_height)
+        else:
+            # No gap or overlap - find actual split position
+            curr_row_y2 = [cells_flat[cell_id]['bbox'][3] for cell_id in row_boundaries[i]['row_cells']]
+            next_row_y1 = [cells_flat[cell_id]['bbox'][1] for cell_id in row_boundaries[i + 1]['row_cells']]
+            
+            max_curr_y2 = max(curr_row_y2)
+            min_next_y1 = min(next_row_y1)
+            
+            # Mark between the actual closest cells
+            if min_next_y1 > max_curr_y2:
+                mark_range(horizontal_gt, max_curr_y2, min_next_y1, orig_height)
+            else:
+                # Overlap - mark fixed width at midpoint
+                split_pos = (max_curr_y2 + min_next_y1) / 2
+                mark_range(horizontal_gt, split_pos - split_width/2, split_pos + split_width/2, orig_height)
+    
+    # 3. Gap from last row bottom to image bottom
+    if row_boundaries:
+        mark_range(horizontal_gt, row_boundaries[-1]['max_y'], orig_height, orig_height)
+    
+    # Mark VERTICAL gaps (between columns)
+    # 1. Gap from image left to first column left
+    if col_boundaries:
+        mark_range(vertical_gt, 0, col_boundaries[0]['min_x'], orig_width)
+    
+    # 2. Gaps between consecutive columns
+    for i in range(len(col_boundaries) - 1):
+        gap_start = col_boundaries[i]['max_x']
+        gap_end = col_boundaries[i + 1]['min_x']
+        
+        if gap_end > gap_start:  # Actual gap exists
+            mark_range(vertical_gt, gap_start, gap_end, orig_width)
+        else:
+            # No gap or overlap - use col_cells to find actual split position
+            curr_col_x2 = [cells_flat[cell_id]['bbox'][2] for cell_id in col_boundaries[i]['col_cells']]
+            next_col_x1 = [cells_flat[cell_id]['bbox'][0] for cell_id in col_boundaries[i + 1]['col_cells']]
+            
+            max_curr_x2 = max(curr_col_x2)
+            min_next_x1 = min(next_col_x1)
+            
+            # Mark between the actual closest cells
+            if min_next_x1 > max_curr_x2:
+                mark_range(vertical_gt, max_curr_x2, min_next_x1, orig_width)
+            else:
+                # Overlap case - mark fixed width at midpoint
+                split_pos = (max_curr_x2 + min_next_x1) / 2
+                mark_range(vertical_gt, split_pos - split_width/2, split_pos + split_width/2, orig_width)
+    
+    # 3. Gap from last column right to image right
+    if col_boundaries:
+        mark_range(vertical_gt, col_boundaries[-1]['max_x'], orig_width, orig_width)
+    
+    return horizontal_gt, vertical_gt
+
+
+class BasicBlock(nn.Module):
+    """Basic ResNet block with halved channels"""
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        
+        self.downsample = None
+        if stride != 1 or inplanes != planes:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes)
+            )
+    
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        
+        out += residual
+        out = self.relu(out)
+        return out
+
+class ModifiedResNet18(nn.Module):
+    """ResNet-18 with removed maxpool and halved channels"""
+    def __init__(self):
+        super().__init__()
+        # First conv block - halved channels: 64→32
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.relu = nn.ReLU(inplace=True)
+        # Skip maxpool - this is the removal mentioned in paper
+        
+        # ResNet layers with halved channels
+        self.layer1 = self._make_layer(32, 32, 2, stride=1)    # Original: 64
+        self.layer2 = self._make_layer(32, 64, 2, stride=2)    # Original: 128
+        self.layer3 = self._make_layer(64, 128, 2, stride=2)   # Original: 256
+        self.layer4 = self._make_layer(128, 256, 2, stride=2)  # Original: 512
+        
+    def _make_layer(self, inplanes, planes, blocks, stride=1):
+        layers = []
+        layers.append(BasicBlock(inplanes, planes, stride))
+        for _ in range(1, blocks):
+            layers.append(BasicBlock(planes, planes))
+        return nn.Sequential(*layers)
+    
+    def forward(self, x):
+        x = self.conv1(x)    # [B, 32, 480, 480]
+        x = self.bn1(x)
+        x = self.relu(x)
+        # No maxpool here - this is the key modification
+        
+        x = self.layer1(x)   # [B, 32, 480, 480]
+        x = self.layer2(x)   # [B, 64, 240, 240]
+        x = self.layer3(x)   # [B, 128, 120, 120]
+        x = self.layer4(x)   # [B, 256, 60, 60]
+        return x
+
+class FPN(nn.Module):
+    """Feature Pyramid Network outputting 128 channels at H/2×W/2"""
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(256, 128, kernel_size=1)
+        
+    def forward(self, x):
+        # x is [B, 256, 60, 60] from ResNet
+        x = self.conv(x)  # [B, 128, 60, 60]
+        # Upsample to H/2×W/2 = 480×480
+        x = F.interpolate(x, size=(480, 480), mode='bilinear', align_corners=False)
+        return x  # [B, 128, 480, 480]
+
+class SplitModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = ModifiedResNet18()
+        self.fpn = FPN()
+        
+        # Learnable weights for global feature averaging
+        self.h_global_weight = nn.Parameter(torch.randn(480))  # For width dimension
+        self.v_global_weight = nn.Parameter(torch.randn(480))  # For height dimension
+        
+        # Local feature processing - reduce to 1 channel then treat spatial as features
+        self.h_local_conv = nn.Conv2d(128, 1, kernel_size=1)
+        self.v_local_conv = nn.Conv2d(128, 1, kernel_size=1)
+        
+        # Fix: Correct feature dimensions - 128 + W/4 = 128 + 120 = 248
+        feature_dim = 128 + 120  # Global + Local features
+
+        # Positional embeddings (1D as mentioned in paper)
+        self.h_pos_embed = nn.Parameter(torch.randn(480, feature_dim))
+        self.v_pos_embed = nn.Parameter(torch.randn(480, feature_dim))
+
+        # Transformers with correct dimensions
+        self.h_transformer = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=feature_dim, nhead=8, dim_feedforward=2048,
+                dropout=0.1, batch_first=True
+            ),
+            num_layers=3
+        )
+        self.v_transformer = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=feature_dim, nhead=8, dim_feedforward=2048,
+                dropout=0.1, batch_first=True
+            ),
+            num_layers=3
+        )
+
+        # Classification heads
+        self.h_classifier = nn.Linear(feature_dim, 1)
+        self.v_classifier = nn.Linear(feature_dim, 1)
+        
+    def forward(self, x):
+        # Input: [B, 3, 960, 960]
+        features = self.backbone(x)      # [B, 256, 60, 60]
+        F_half = self.fpn(features)      # [B, 128, 480, 480] - This is F1/2
+
+        B, C, H, W = F_half.shape        # B, 128, 480, 480
+
+        # HORIZONTAL FEATURES (for row splitting)
+        # Global: learnable weighted average along width dimension
+        F_RG = torch.einsum('bchw,w->bch', F_half, self.h_global_weight)  # [B, 128, 480]
+        F_RG = F_RG.transpose(1, 2)  # [B, 480, 128]
+
+        # Local: 1×4 AvgPool to get 120 features (W/4), then 1×1 conv to 1 channel
+        F_RL_pooled = F.avg_pool2d(F_half, kernel_size=(1, 4))  # [B, 128, 480, 120]
+        F_RL = self.h_local_conv(F_RL_pooled)  # [B, 1, 480, 120]
+        F_RL = F_RL.squeeze(1)  # [B, 480, 120] - spatial becomes features
+
+        # Concatenate: [B, 480, 128+120=248]
+        F_RG_L = torch.cat([F_RG, F_RL], dim=2)
+
+        # Add positional embeddings
+        F_RG_L = F_RG_L + self.h_pos_embed
+
+        # VERTICAL FEATURES (for column splitting)
+        # Global: learnable weighted average along height dimension
+        F_CG = torch.einsum('bchw,h->bcw', F_half, self.v_global_weight)  # [B, 128, 480]
+        F_CG = F_CG.transpose(1, 2)  # [B, 480, 128]
+
+        # Local: 4×1 AvgPool to get 120 features (H/4), then 1×1 conv to 1 channel
+        F_CL_pooled = F.avg_pool2d(F_half, kernel_size=(4, 1))  # [B, 128, 120, 480]
+        F_CL = self.v_local_conv(F_CL_pooled)  # [B, 1, 120, 480]
+        F_CL = F_CL.squeeze(1)  # [B, 120, 480]
+        F_CL = F_CL.transpose(1, 2)  # [B, 480, 120] - transpose to get spatial as features
+
+        # Concatenate: [B, 480, 128+120=248]
+        F_CG_L = torch.cat([F_CG, F_CL], dim=2)
+
+        # Add positional embeddings
+        F_CG_L = F_CG_L + self.v_pos_embed
+
+        # Transformer processing
+        F_R = self.h_transformer(F_RG_L)  # [B, 480, 368]
+        F_C = self.v_transformer(F_CG_L)  # [B, 480, 368]
+
+        # Binary classification at 480 resolution
+        h_logits = self.h_classifier(F_R).squeeze(-1)  # [B, 480]
+        v_logits = self.v_classifier(F_C).squeeze(-1)  # [B, 480]
+
+        # return at 480 resolution (upsample happens AFTER loss computation)
+        return torch.sigmoid(h_logits), torch.sigmoid(v_logits)  # [B, 480]
+
+def focal_loss(predictions, targets, alpha=1.0, gamma=2.0):
+    """Focal loss exactly as specified in paper"""
+    ce_loss = F.binary_cross_entropy(predictions, targets, reduction='none')
+    pt = torch.where(targets == 1, predictions, 1 - predictions)
+    focal_weight = alpha * (1 - pt) ** gamma
+    return (focal_weight * ce_loss).mean()
+
+def post_process_predictions(h_pred, v_pred, threshold=0.5):
+    """
+    Simple post-processing to convert predictions to binary masks
+    """
+    h_binary = (h_pred > threshold).float()
+    v_binary = (v_pred > threshold).float()
+
+    return h_binary, v_binary
+
+class TableDataset(Dataset):
+    def __init__(self, hf_dataset):
+        self.hf_dataset = hf_dataset
+        self.transform = transforms.Compose([
+            transforms.Resize((960, 960)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    def __len__(self):
+        return len(self.hf_dataset)
+
+    def __getitem__(self, idx):
+        item = self.hf_dataset[idx]
+
+        image = item['image'].convert('RGB')
+        image_transformed = self.transform(image)
+
+        # generate GT at 960 resolution
+        h_gt_960, v_gt_960 = get_ground_truth_auto_gap(
+            item['image'],  # original PIL image for dimensions
+            item['cells'],
+            item['otsl'],
+        )
+
+        # downsample to 480 for loss computation (take every 2nd element)
+        h_gt_480 = [h_gt_960[i] for i in range(0, 960, 2)]  # [480]
+        v_gt_480 = [v_gt_960[i] for i in range(0, 960, 2)]  # [480]
+
+        return (
+            image_transformed,
+            torch.tensor(h_gt_480, dtype=torch.float),  # [480] for training loss
+            torch.tensor(v_gt_480, dtype=torch.float),  # [480] for training loss
+            torch.tensor(h_gt_960, dtype=torch.float),  # [960] for metrics
+            torch.tensor(v_gt_960, dtype=torch.float),  # [960] for metrics
+        )
+
+