{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"metadata":{},"cell_type":"markdown","source":"# Train notebook\n\nTrain notebook for attending Global Wheat Detection competition.\n\nRelated links:\n\nNotebook used for result submission, can be found — [here](https://www.kaggle.com/code/pavlozelinskiy/global-wheat-detection-submit-v1).\n\nBest model dataset, can be found — [here](https://www.kaggle.com/datasets/pavlozelinskiy/global-wheat-detection-weights-v1).\n"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:04:48.516427900Z","start_time":"2026-09-24T17:04:48.464681Z"}},"cell_type":"code","source":"%matplotlib inline\n\n# region Imports\nimport numpy as np\nimport pandas as pd\nimport torch\nimport timm\nimport torch.nn as nn\n\nimport os\nimport math\nimport time\nfrom pathlib import Path\n# from kaggle_second_competition.dataset import ImageDataset\n# from kaggle_second_competition.detector import build_detector, describe_detector\n# from kaggle_second_competition.detection_metric import competition_metric, images_average_precision, \\\n#     match_counts\n# from kaggle_second_competition.detection_transform import build_train_transform, \\\n#     build_validation_transform, detection_collate\n\nfrom torchvision import transforms as T\nfrom torchvision.ops import box_iou\nfrom torchvision.transforms import v2\nfrom torchvision import tv_tensors\n\nfrom PIL import Image\nfrom PIL.Image import Resampling\nfrom torch.utils.data import DataLoader\n\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\n\nfrom sklearn.model_selection import StratifiedGroupKFold\n\n# endregion\n\n# region Constants\nNUM_CLASSES = 2  # Note: one for background, one for wheat label\nDEVICE = 'cuda'  # 12Gb VRam\nEPOCHS = 10\nORIGINAL_IMAGE_SIZE = 1024\n\nMIN_VISIBILITY = 0.3  # TODO: needs more tuning\n\nUNKNOWN_SOURCE = 'unknown'\nORIGINAL_COLUMNS = ['image_id', 'width', 'height', 'bbox', 'source']\n\n# region Data analysis constants\nMAX_BOX_LIMIT_CHECK = 100  # Note: 100, is not a hard limit\nTOO_BIG_LONGEST_SIDE = ORIGINAL_IMAGE_SIZE / 2  # Note: I don't think that wheat head half of the image size is valid\nTOO_SMALL_SIDE = 10  # Note: area < 100 px² is probably too low, for 1024x1024 image\nHIGH_OVERLAP_IOU = 0.5  # Note: High overlap does not consider as outlier but rather to compare against duplicate IoU\nDUPLICATE_IOU = 0.9\n# endregion\n\n# endregion\n\n# region Configuration\nRUN_MODEL_TRAINING = True\n# endregion\n\n# region Path related\nIS_KAGGLE = bool(os.environ.get('KAGGLE_KERNEL_RUN_TYPE'))\nKAGGLE_DIR = Path('/kaggle') if IS_KAGGLE else Path.cwd().parent / 'kaggle'\n\nINPUT_DIR = KAGGLE_DIR / 'input' / 'competitions' / 'global-wheat-detection'\nOUTPUT_DIR = KAGGLE_DIR / 'working'\n\nINSPECT_ONLY_TRAIN_CSV_PATH = OUTPUT_DIR / 'inspect_only_cleaned_train.csv'\n\nMODEL_DIR = OUTPUT_DIR\nBEST_MODEL_PATH = MODEL_DIR / 'best_model.pth'\n\nfor label, d in [('INPUT_DIR', INPUT_DIR), ('MODEL_DIR', MODEL_DIR)]:\n    if not d.is_dir():\n        raise FileNotFoundError(\n            f\"{label} missing: {d}\\nmounted: {[p.name for p in Path('/kaggle/input').iterdir()]}\"\n            if IS_KAGGLE else f\"{label} missing: {d}\"\n        )\n# endregion","outputs":[],"execution_count":12},{"metadata":{},"cell_type":"markdown","source":"# Define class for PyTorch Dataset\nto run training locally directly on Windows .py version of this class is used"},{"metadata":{},"cell_type":"code","outputs":[],"execution_count":null,"source":"# Copy pasted from dataset.py and detection_transform.py and detector.py and detection_metric.py\n\nfrom pathlib import Path\n\nimport numpy as np\nimport torch\nfrom PIL import Image\nfrom torch.utils.data import Dataset\n\nWHEAT_HEAD_CLASS_ID = 1\n\n\nclass ImageDataset(Dataset):\n    \"\"\"\n    One item == one image plus all of its boxes.\n        transform: callable(PIL.Image, ndarray[N, 4]) -> (image_tensor, boxes[M, 4])\n    Boxes come out of here as absolute-pixel xyxy (x_min, y_min, x_max, y_max).\n    \"\"\"\n\n    def __init__(self, data_frame, image_dir, transform, has_label=True):\n        self.image_dir = Path(image_dir)\n        self.transform = transform\n\n        if not has_label:\n            self.image_ids = [str(i) for i in data_frame['image_id']]\n            self.boxes = None\n            self.source = [None] * len(self.image_ids)\n            return\n\n        groups = data_frame.groupby('image_id', sort=False).indices\n        boxes = self.parse_bounding_boxes(data_frame['bbox'])\n        is_finite_mask = np.isfinite(boxes).all(axis=1)\n        sources = data_frame['source'].to_numpy()\n\n        # Note: wrapping into `str` to prevent TypeError when working with path\n        self.image_ids = [str(i) for i in groups]\n        self.boxes = [boxes[rows[is_finite_mask[rows]]] for rows in groups.values()]\n        self.source = [sources[rows[0]] for rows in groups.values()]\n\n    @staticmethod\n    def parse_bounding_boxes(bounding_box_column):\n        \"\"\"\n        Convert `xywh` style bounding boxes to `xyxy`.\n        \"\"\"\n        xyxy = np.full((len(bounding_box_column), 4), np.nan, dtype='float32')\n        is_box = bounding_box_column.notna().to_numpy()\n        if not is_box.any():\n            return xyxy\n\n        xywh = (bounding_box_column[is_box]\n                .str.strip('[]')\n                .str.split(',', expand=True).to_numpy(dtype='float32'))\n        assert xywh.shape[1] == 4, f'expected 4 numbers per bbox, parsed {xywh.shape[1]}'\n\n        xyxy[is_box, :2] = xywh[:, :2]\n        xyxy[is_box, 2:] = xywh[:, :2] + xywh[:, 2:]\n        return xyxy\n\n    def __len__(self):\n        return len(self.image_ids)\n\n    def __getitem__(self, i):\n        image_id = self.image_ids[i]\n        # Note: .convert('RGB') added as a sanity check in hidden test set\n        img = Image.open(self.image_dir / f'{image_id}.jpg').convert('RGB')\n        boxes = self.boxes[i] if self.boxes is not None else np.zeros((0, 4), dtype='float32')\n\n        img, boxes = self.transform(img, boxes)\n\n        target = {'image_id': image_id, 'source': self.source[i]}\n        if self.boxes is not None:\n            target['boxes'] = torch.as_tensor(np.asarray(boxes, dtype='float32')).reshape(-1, 4)\n            target['labels'] = torch.full((len(target['boxes']),), WHEAT_HEAD_CLASS_ID, dtype=torch.int64)\n        return img, target\n\n\nimport numpy as np\nimport torch\nfrom torchvision import tv_tensors\nfrom torchvision.transforms import v2\n\n\ndef box_areas(boxes):\n    widths = (boxes[:, 2] - boxes[:, 0]).clamp(min=0)\n    heights = (boxes[:, 3] - boxes[:, 1]).clamp(min=0)\n    return widths * heights\n\n\nclass BoxAwareTransform:\n    \"\"\"\n    Box-aware transform satisfying ImageDataset's contract:\n        callable(PIL.Image, ndarray[N, 4] xyxy) -> (float tensor[3, H, W], ndarray[M, 4] xyxy)\n    \"\"\"\n\n    def __init__(self, pipeline, min_visibility=0.0):\n        self.pipeline = pipeline\n        self.clamp_boxes = v2.ClampBoundingBoxes(clamping_mode='hard')\n        self.to_float = v2.ToDtype(torch.float32, scale=True)\n        self.min_visibility = min_visibility\n\n    def __call__(self, image, boxes):\n        # Note: most of this logic is unused before we have some Crop, or Scale augmentations\n        #  clamping to < 1px size, is important when resize from 1024 to lower resolution will be done.\n        boxes = tv_tensors.BoundingBoxes(\n            data=torch.as_tensor(np.asarray(boxes, dtype='float32')).reshape(-1, 4),\n            format='XYXY', canvas_size=(image.height, image.width), clamping_mode=None)\n\n        image, boxes = self.pipeline(v2.functional.to_image(image), boxes)\n\n        clamped = self.clamp_boxes(boxes).as_subclass(torch.Tensor)\n        widths = clamped[:, 2] - clamped[:, 0]\n        heights = clamped[:, 3] - clamped[:, 1]\n        visibility = (widths * heights) / box_areas(boxes.as_subclass(torch.Tensor)).clamp(min=1e-6)\n        keep = (visibility >= self.min_visibility) & (widths >= 1) & (heights >= 1)\n\n        # Note: float32 scaled to [0, 1], bc the detector normalizes it\n        image = self.to_float(image)\n        return image, clamped[keep].numpy()\n\n\ndef build_train_transform(image_size, min_visibility):\n    steps = [\n        v2.Resize(size=(image_size, image_size)),\n        v2.RandomHorizontalFlip(p=0.5),\n        v2.RandomVerticalFlip(p=0.5),\n        v2.RandomApply([v2.RandomRotation(degrees=(90, 90))], p=0.5),\n        v2.ColorJitter(brightness=(0.7, 1.3), contrast=(0.8, 1.2), saturation=(0.7, 1.3), hue=(-0.08, 0.08)),\n        v2.RandomGrayscale(p=0.05),\n        v2.RandomApply([v2.GaussianBlur(kernel_size=9, sigma=(0.5, 2.0))], p=0.1),\n    ]\n    return BoxAwareTransform(pipeline=v2.Compose(steps), min_visibility=min_visibility)\n\n\ndef build_validation_transform(image_size):\n    steps = [v2.Resize(size=(image_size, image_size))]\n    return BoxAwareTransform(v2.Compose(steps), min_visibility=0.0)\n\n\ndef detection_collate(batch):\n    \"\"\"\n    The default collate cannot stack targets, because the number of boxes differs per image.\n    Returns (list[Tensor], list[dict]), which is exactly what torchvision's detectors accept.\n    \"\"\"\n    images, targets = zip(*batch)\n    shapes = {tuple(image.shape) for image in images}\n    assert len(shapes) == 1, f'images in one batch must share a shape, got {shapes}'\n    return list(images), list(targets)\n\nfrom functools import partial\n\nimport torch.nn as nn\nfrom torchvision.models.detection import (fasterrcnn_resnet50_fpn_v2, fcos_resnet50_fpn,\n                                          retinanet_resnet50_fpn_v2)\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.fcos import FCOSClassificationHead\nfrom torchvision.models.detection.retinanet import RetinaNetClassificationHead\n\nDETECTION_HEADS = ('faster_rcnn', 'retinanet', 'fcos')\n\n\ndef build_detector(head, image_size, num_classes, max_detections_per_image=300,\n                   score_threshold=0.05, nms_threshold=0.5, pretrained=True):\n    assert head in DETECTION_HEADS, f'unknown head {head!r}, expected one of {DETECTION_HEADS}'\n    weights = 'DEFAULT' if pretrained else None\n    shared = dict(weights=weights, min_size=image_size, max_size=image_size)\n    if not pretrained:\n        shared['weights_backbone'] = None\n\n    if head == 'faster_rcnn':\n        model = fasterrcnn_resnet50_fpn_v2(\n            box_detections_per_img=max_detections_per_image,\n            box_score_thresh=score_threshold, box_nms_thresh=nms_threshold, **shared)\n        # The COCO predictor emits 91 classes; only this last layer is replaced.\n        in_features = model.roi_heads.box_predictor.cls_score.in_features\n        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n        return model\n\n    one_stage = dict(detections_per_img=max_detections_per_image,\n                     score_thresh=score_threshold, nms_thresh=nms_threshold, **shared)\n\n    if head == 'retinanet':\n        model = retinanet_resnet50_fpn_v2(**one_stage)\n        # norm_layer matches what the _v2 builder uses for its own head, so the replacement is\n        # structurally identical to the one it shipped with\n        model.head.classification_head = RetinaNetClassificationHead(\n            in_channels=model.backbone.out_channels,\n            num_anchors=model.anchor_generator.num_anchors_per_location()[0],\n            num_classes=num_classes,\n            norm_layer=partial(nn.GroupNorm, 32))\n        return model\n\n    model = fcos_resnet50_fpn(**one_stage)\n    model.head.classification_head = FCOSClassificationHead(\n        in_channels=model.backbone.out_channels,\n        num_anchors=model.anchor_generator.num_anchors_per_location()[0],\n        num_classes=num_classes)\n    return model\n\n\nimport numpy as np\nimport torch\nfrom torchvision.ops import box_iou\n\n\"\"\"\nTP if overlap reaches the threshold, a FP otherwise.\nGround-truth boxes nobody claimed are FN.\nAs per competitions:\n    score = TP / (TP + FP + FN)\n    or\n    TP + FP + FN == n_predictions + n_ground_truth - TP\n\nThe image's score is the mean over the six thresholds.\n\"\"\"\n\n# TODO: Try to use torchmetrics in future instead of writing own metrics implementation\nEMPTY_IMAGE_SCORE = 1.0  # Note: as per competition rules\n\ndef competition_metric(predictions, targets, confidence_threshold, thresholds):\n    precisions = images_average_precision(predictions, targets, confidence_threshold, thresholds)\n    return float(np.mean(precisions)) if precisions else 0.0\n\n\ndef sorted_iou(predicted_boxes, ground_truth_boxes, confidences):\n    \"\"\"IoU matrix with the prediction rows put in descending confidence order.\n\n    Returns (iou, n_predictions, n_ground_truth). iou is None when either side is empty - there\n    is nothing to match then, and what that means differs per caller.\n    \"\"\"\n    predicted_boxes = torch.as_tensor(predicted_boxes, dtype=torch.float32).detach().cpu().reshape(-1, 4)\n    ground_truth_boxes = torch.as_tensor(ground_truth_boxes, dtype=torch.float32).detach().cpu().reshape(-1, 4)\n    n_predictions, n_ground_truth = len(predicted_boxes), len(ground_truth_boxes)\n\n    if n_predictions == 0 or n_ground_truth == 0:\n        return None, n_predictions, n_ground_truth\n\n    if confidences is not None:\n        confidences = torch.as_tensor(confidences, dtype=torch.float32).detach().cpu().reshape(-1)\n        assert len(confidences) == n_predictions, f'{len(confidences)} confidences for {n_predictions} boxes'\n        predicted_boxes = predicted_boxes[torch.argsort(confidences, descending=True)]\n\n    return box_iou(predicted_boxes, ground_truth_boxes).numpy(), n_predictions, n_ground_truth\n\n\ndef filter_by_confidence(prediction, confidence_threshold):\n    boxes, confidences = prediction['boxes'], prediction['scores']\n    if confidence_threshold is None:\n        return boxes, confidences\n    keep = confidences >= confidence_threshold\n    return boxes[keep], confidences[keep]\n\n\ndef image_average_precision(predicted_boxes, ground_truth_boxes, confidences, thresholds):\n    iou, n_predictions, n_ground_truth = sorted_iou(predicted_boxes, ground_truth_boxes, confidences)\n    if iou is None:\n        # Nothing predicted on a real image is all FN; anything predicted on an empty one is all\n        # FP. Nothing to find and nothing predicted means nothing was got wrong.\n        return EMPTY_IMAGE_SCORE if n_predictions == 0 and n_ground_truth == 0 else 0.0\n\n    total = 0.0\n    for threshold in thresholds:\n        matches = count_matches(iou, threshold)\n        total += matches / (n_predictions + n_ground_truth - matches)\n    return total / len(thresholds)\n\n\ndef match_counts(predictions, targets, confidence_threshold, iou_threshold):\n    \"\"\"\n    (Diagnostic only) Total (TP, FP, FN) over every image, at one confidence and one IoU threshold.\n    \"\"\"\n    assert len(predictions) == len(targets), f'{len(predictions)} predictions, {len(targets)} targets'\n    true_positives = false_positives = false_negatives = 0\n    for prediction, target in zip(predictions, targets):\n        boxes, confidences = filter_by_confidence(prediction, confidence_threshold)\n        iou, n_predictions, n_ground_truth = sorted_iou(boxes, target['boxes'], confidences)\n        matches = count_matches(iou, iou_threshold) if iou is not None else 0\n        true_positives += matches\n        false_positives += n_predictions - matches\n        false_negatives += n_ground_truth - matches\n    return true_positives, false_positives, false_negatives\n\n\ndef images_average_precision(predictions, targets, confidence_threshold, thresholds):\n    \"\"\"\n    Per-image scores, in input order.\n    \"\"\"\n    assert len(predictions) == len(targets), f'{len(predictions)} predictions, {len(targets)} targets'\n    scored = []\n    for prediction, target in zip(predictions, targets):\n        boxes, confidences = filter_by_confidence(prediction, confidence_threshold)\n        scored.append(image_average_precision(boxes, target['boxes'], confidences, thresholds))\n    return scored\n\n\ndef count_matches(iou, threshold):\n    \"\"\"\n    IoU: ndarray[n_predictions, n_ground_truth], rows already in descending confidence order.\n    (as per competition rules)\n    Greedy: each prediction in turn claims the best ground-truth box still unclaimed.\n    \"\"\"\n    claimed = np.zeros(iou.shape[1], dtype=bool)\n    matches = 0\n    for row in iou:\n        candidates = np.where(claimed, -1.0, row)\n        best = int(candidates.argmax())\n        if candidates[best] >= threshold:\n            claimed[best] = True\n            matches += 1\n    return matches\n"},{"metadata":{},"cell_type":"markdown","source":"# Common functions"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:04:55.040298200Z","start_time":"2026-09-24T17:04:53.976079700Z"}},"cell_type":"code","source":"def init_data_frame():\n    return pd.read_csv(f'{INPUT_DIR}/train.csv')\n\n\ndef init_bounding_boxes():\n    return ImageDataset.parse_bounding_boxes(data_frame['bbox'])\n\n\ndef calculate_bounding_box_area():\n    widths = bounding_boxes[:, 2] - bounding_boxes[:, 0]\n    heights = bounding_boxes[:, 3] - bounding_boxes[:, 1]\n    data_frame['area'] = (widths * heights).astype('float64')\n    data_frame['area_geometric_mean'] = np.sqrt(data_frame['area'])\n\n\ndef calculate_longest_side():\n    widths = bounding_boxes[:, 2] - bounding_boxes[:, 0]\n    heights = bounding_boxes[:, 3] - bounding_boxes[:, 1]\n    data_frame['longest'] = np.maximum(widths, heights)\n\n\ndef calculate_max_iou():\n    \"\"\"\n    Compute IoU for every box it has with any other box in the current image\n    \"\"\"\n    max_iou = np.zeros(len(data_frame))\n    rows_indices_by_image = init_rows_indices_by_image()\n    for rows in rows_indices_by_image.values():\n        if len(rows) > 1:\n            boxes = torch.from_numpy(bounding_boxes[rows])\n            max_iou[rows] = box_iou(boxes, boxes).fill_diagonal_(0).max(dim=1).values.numpy()\n    data_frame['max_iou'] = max_iou\n\n\ndef init_all_image_ids():\n    \"\"\"\n    Retrieving from image directory, bc images with no bounding boxes is not presented in train.csv file\n    \"\"\"\n    all_image_ids = {p.stem for p in (INPUT_DIR / 'train').glob('*.jpg')}  # Note: This is Set, if you forget\n    unique_with_boxes_image_ids = set(data_frame['image_id'])\n    assert unique_with_boxes_image_ids <= all_image_ids, 'train.csv references images missing from disk'\n    assert data_frame.groupby('image_id')['source'].nunique().max() == 1, 'an image has boxes from 2+ sources'\n    return all_image_ids\n\n\ndef init_rows_indices_by_image():\n    return data_frame.groupby('image_id').indices\n\n\ndata_frame = init_data_frame()\nbounding_boxes = init_bounding_boxes()\ncalculate_bounding_box_area()\ncalculate_longest_side()\ncalculate_max_iou()\n\nall_image_ids = init_all_image_ids()\nrows_indices_by_image = init_rows_indices_by_image()\n","outputs":[],"execution_count":13},{"metadata":{},"cell_type":"markdown","source":"# Data analysis"},{"metadata":{},"cell_type":"markdown","source":"## Dataset visualization"},{"metadata":{},"cell_type":"code","outputs":[],"execution_count":null,"source":"SOURCES_TO_SHOW = ('arvalis_1', 'ethz_1', 'inrae_1', 'usask_1')\nALL_SOURCES = ('arvalis_1', 'arvalis_2', 'arvalis_3', 'ethz_1', 'inrae_1', 'usask_1', 'rres_1')\n\n\ndef visualize_images_from_different_sources(sources=SOURCES_TO_SHOW, display_size=512, seed=42):\n    rng = np.random.default_rng(seed)\n    fig, axes = plt.subplots(nrows=1, ncols=len(sources), figsize=(4.5 * len(sources), 5.2), squeeze=False)\n\n    for ax, source in zip(axes[0], sources):\n        source_mask = data_frame['source'] == source\n        source_rows = data_frame[source_mask]\n        image_id = rng.choice(source_rows['image_id'].unique())\n        image_rows = source_rows[source_rows['image_id'] == image_id]\n\n        bounding_boxes = ImageDataset.parse_bounding_boxes(image_rows['bbox'])\n\n        image = Image.open(f'{INPUT_DIR}/train/{image_id}.jpg')\n        scale = display_size / image.width\n        image = image.resize(size=(display_size, round(image.height * scale)), resample=Resampling.BILINEAR)\n        bounding_boxes = bounding_boxes * scale\n\n        ax.imshow(image)\n        for x1, y1, x2, y2 in bounding_boxes:\n            ax.add_patch(\n                patches.Rectangle(xy=(x1, y1), width=x2 - x1, height=y2 - y1, linewidth=1.0, edgecolor='violet',\n                                  facecolor='none'))\n\n        widths = bounding_boxes[:, 2] - bounding_boxes[:, 0]\n        heights = bounding_boxes[:, 3] - bounding_boxes[:, 1]\n        areas = widths * heights\n        scaled_geometric_means = np.sqrt(areas) / scale\n\n        title_first_line = f'{source}  |  {len(bounding_boxes)} boxes'\n        title_second_line = f'{image_id}  |  median box {np.median(scaled_geometric_means):.0f}px'\n        ax.set_title(f'{title_first_line}\\n{title_second_line}', fontsize=9)\n        ax.axis('off')\n\n    plt.tight_layout()\n    plt.show()\n\n\nvisualize_images_from_different_sources()\n"},{"metadata":{},"cell_type":"markdown","source":"## Exploratory data analysis"},{"metadata":{},"cell_type":"code","outputs":[],"execution_count":null,"source":"all_image_ids = init_all_image_ids()\nunique_with_boxes_image_ids = set(data_frame['image_id'])\nempty_image_ids = sorted(all_image_ids - unique_with_boxes_image_ids)\n\nboxes_per_image = data_frame.groupby('image_id')['image_id'].count().reindex(sorted(all_image_ids), fill_value=0)\n\nprint('=== Whole dataset ===')\nprint(f'training images  {len(all_image_ids)} '\n      f'({len(unique_with_boxes_image_ids)} in train.csv + {len(empty_image_ids)} with no boxes)')\nprint(f'total boxes      {len(data_frame)}')\nprint(f'mean boxes/image {boxes_per_image.mean():.2f} (images without boxes counted as 0)')\nprint(\n    f'median side      {data_frame['area_geometric_mean'].median():.0f}px (median box area {data_frame[\"area\"].median():.0f} px²)')\n\nper_image = data_frame.groupby('image_id').agg(source=('source', 'first'), boxes=('bbox', 'size'))\nper_source = pd.DataFrame({\n    'images': per_image.groupby('source').size(),\n    '% of all imgs': per_image.groupby('source').size() / len(all_image_ids) * 100,\n    'mean boxes/img': per_image.groupby('source')['boxes'].mean(),\n    'median side px': data_frame.groupby('source')['area_geometric_mean'].median(),\n    'median area px²': data_frame.groupby('source')['area'].median(),\n}).sort_values('images', ascending=False)\n\nprint(f'\\n=== Per source ({len(per_source)} sources) ===')\nprint(per_source.round(1).to_string())\n\nprint('\\n=== Images with no boxes ===')\nprint(f'{len(empty_image_ids)} images ({len(empty_image_ids) / len(all_image_ids):.1%} of all), '\n      f'e.g. {empty_image_ids[:4]}')\nprint('Warning: \"source\" for such images is not defined!')"},{"metadata":{},"cell_type":"markdown","source":"## Outliers"},{"metadata":{},"cell_type":"code","outputs":[],"execution_count":null,"source":"def get_outlier_entry(mask):\n    return sorted(data_frame.loc[mask, 'image_id'].unique()), mask\n\n\ndata_frame['on_border'] = ((bounding_boxes[:, 0] <= 0) | (bounding_boxes[:, 1] <= 0) |\n                           (bounding_boxes[:, 2] >= ORIGINAL_IMAGE_SIZE) | (\n                                   bounding_boxes[:, 3] >= ORIGINAL_IMAGE_SIZE))\n\nempty_box_image_ids = sorted(all_image_ids - unique_with_boxes_image_ids)\nmax_box_limit_image_ids = sorted(boxes_per_image.index[boxes_per_image >= MAX_BOX_LIMIT_CHECK])\n\ntoo_big_box_mask = get_outlier_entry(mask=data_frame['longest'] > TOO_BIG_LONGEST_SIDE)\ntoo_small_box_mask = get_outlier_entry(mask=data_frame['area_geometric_mean'] < TOO_SMALL_SIDE)\noverlap_box_mask = get_outlier_entry(\n    mask=data_frame['max_iou'].between(HIGH_OVERLAP_IOU, DUPLICATE_IOU, inclusive='right'))\nduplicate_box_mask = get_outlier_entry(mask=data_frame['max_iou'] > DUPLICATE_IOU)\noutliers = {\n    'no boxes': (empty_box_image_ids, None),\n    f'{MAX_BOX_LIMIT_CHECK}+ boxes': (max_box_limit_image_ids, None),\n    f'too big (longest side > {TOO_BIG_LONGEST_SIDE:.0f}px)': too_big_box_mask,\n    f'too small (side < {TOO_SMALL_SIDE}px)': too_small_box_mask,\n    f'overlap (IoU {HIGH_OVERLAP_IOU}..{DUPLICATE_IOU})': overlap_box_mask,\n    f'duplicate (IoU > {DUPLICATE_IOU})': duplicate_box_mask,\n}\n\n\ndef prepare_outliers() -> pd.DataFrame:\n    outliers_summary = []\n    for name, (image_ids, mask) in outliers.items():\n        box_count = int(mask.sum()) if mask is not None else '-'\n        border_count = int((mask & data_frame['on_border']).sum()) if mask is not None else '-'\n\n        outliers_summary.append({\n            'category': name,\n            'images': len(image_ids),\n            'boxes': box_count,\n            'of which on border': border_count\n        })\n\n    return pd.DataFrame(outliers_summary)\n\n\noutliers_summary = prepare_outliers()\n\noutliers_summary_by_source = {\n    name: pd.Series(image_ids, dtype=object).map(per_image['source']).fillna('unknown').value_counts()\n    for name, (image_ids, _) in outliers.items()\n}\noutliers_summary_by_source = pd.DataFrame(outliers_summary_by_source)\n\nprint('=== Outliers ===')\nprint(outliers_summary.to_string(index=False))\n\nprint('\\n=== Outlier images per source ===')\nprint(outliers_summary_by_source.fillna(0).astype(int).T.to_string())\n\n\n# Note: marks flagged boxes with thick red line and circle in the center\ndef plot_outliers(outliers, display_size=512, seed=42, n_cols=3):\n    rng = np.random.default_rng(seed)\n    n_rows = math.ceil(len(outliers) / n_cols)\n    fig, axes = plt.subplots(nrows=n_rows, ncols=n_cols, figsize=(5 * n_cols, 5.4 * n_rows), squeeze=False)\n    for ax in axes.flat:\n        ax.axis('off')\n\n    for ax, (name, (image_ids, mask)) in zip(axes.flat, outliers.items()):\n        if not image_ids:\n            ax.set_title(f'{name}\\n(no images)', fontsize=9)\n            continue\n        image_id = rng.choice(image_ids)\n        rows = rows_indices_by_image.get(image_id, np.array([], dtype=int))\n        flagged_boxes = mask.to_numpy()[rows] if mask is not None else np.zeros(len(rows), dtype=bool)\n\n        image = Image.open(f'{INPUT_DIR}/train/{image_id}.jpg')\n        scale = display_size / image.width\n        image = image.resize((display_size, round(image.height * scale)), resample=Resampling.BILINEAR)\n        ax.imshow(image)\n\n        for (x1, y1, x2, y2), flagged in sorted(zip(bounding_boxes[rows] * scale, flagged_boxes),\n                                                key=lambda box: box[1]):\n            ax.add_patch(\n                patches.Rectangle(xy=(x1, y1), width=x2 - x1, height=y2 - y1, linewidth=2.0 if flagged else 0.7,\n                                  edgecolor='red' if flagged else 'violet', facecolor='none'))\n            if flagged:\n                ax.plot((x1 + x2) / 2, (y1 + y2) / 2, marker='o', markersize=18, mfc='none', mec='red', mew=1.5,\n                        scalex=False, scaley=False)  # Note: disable scale to prevent image from padding\n\n        source = per_image['source'].get(image_id, default='unknown')\n        title_first_line = f'{name}  |  {len(image_ids)} images'\n        title_second_line = f'{image_id}  |  {source}  |  {len(rows)} boxes, {flagged_boxes.sum()} flagged'\n        ax.set_title(f'{title_first_line}\\n{title_second_line}', fontsize=9)\n\n    plt.tight_layout()\n    plt.show()\n\n\nplot_outliers(outliers)\n"},{"metadata":{},"cell_type":"markdown","source":"## Augmentation tester"},{"metadata":{},"cell_type":"code","outputs":[],"execution_count":null,"source":"# TODO: Consider adding mosaic augmentation too?\n# Note: Uses edge values for augmentation showcase\nTEST_AUGMENTATIONS = {\n    f'HorizontalFlip': v2.RandomHorizontalFlip(p=1.0),\n    f'VerticalFlip': v2.RandomVerticalFlip(p=1.0),\n    f'Rotation 90°': v2.RandomRotation(degrees=(90, 90)),\n    # Note: do NOT use\n    f'Rotation 15°': v2.RandomRotation(degrees=(15, 15)),\n    # Note: All ColorJitter looks viable\n    f'ColorJitter: brightness 0.6, contrast 0.7': v2.ColorJitter(brightness=(0.6, 0.6), contrast=(0.7, 0.7)),\n    f'ColorJitter: brightness 1.3, saturation 1.5': v2.ColorJitter(brightness=(1.3, 1.3), saturation=(1.5, 1.5)),\n    f'ColorJitter: hue +0.08': v2.ColorJitter(hue=(0.08, 0.08)),\n    # f'ColorJitter: all combined': v2.ColorJitter(brightness=(1.2, 1.2), saturation=(1.2, 1.2), hue=(0.08, 0.08)),\n    f'ColorJitter: all combined': v2.ColorJitter(brightness=(1.0, 1.0), contrast=(1.2, 1.2), hue=(0.08, 0.08)),\n    f'Grayscale': v2.Grayscale(num_output_channels=3),\n    # Note: look in full res, via double-click on graph\n    f'GaussianBlur: sigma 2': v2.GaussianBlur(kernel_size=19, sigma=(2.0, 2.0)),\n    # Note: TODO: don't see any noticeable effect, consider if it makes sense to use\n    f'Jpeg compression: 50': v2.JPEG(quality=(50, 50)),\n    # Note: do NOT use\n    f'RandomCrop 640': v2.RandomCrop(size=640),\n    f'Resize 512': v2.Resize(size=(512, 512)),\n    # Note: Needs tuning\n    f'Zoom in x2 (RandomResizedCrop)': v2.RandomResizedCrop(size=(1024, 1024), scale=(0.25, 0.25), ratio=(1.0, 1.0)),\n    f'Zoom out x1.5 (RandomZoomOut + Resize)': v2.Compose(  # Note: do NOT use (leaves black borders)\n        [v2.RandomZoomOut(fill=0, side_range=(1.5, 1.5), p=1.0), v2.Resize(size=(1024, 1024))]\n    ),\n}\n\n\ndef box_area(boxes):\n    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])\n\n\ndef augment(transform, image, bboxes):\n    boxes = tv_tensors.BoundingBoxes(torch.from_numpy(bboxes), format='XYXY', canvas_size=(image.height, image.width),\n                                     clamping_mode=None)\n    image, boxes = transform(image, boxes)\n    clamped = v2.ClampBoundingBoxes(clamping_mode='hard')(boxes)\n    return image, clamped.numpy(), box_area(clamped.numpy()) / box_area(boxes.numpy())\n\n\ndef draw_panel(ax, image, boxes, colors, max_dimension, title):\n    ax.imshow(image)\n    for (x1, y1, x2, y2), color in zip(boxes, colors):\n        ax.add_patch(patches.Rectangle(xy=(x1, y1), width=x2 - x1, height=y2 - y1, linewidth=1.0, edgecolor=color,\n                                       facecolor='none', linestyle='--' if color == 'red' else '-'))\n    ax.set_xlim(-0.5, max_dimension - 0.5)\n    ax.set_ylim(max_dimension - 0.5, -0.5)\n    ax.set_title(title, fontsize=9)\n    ax.axis('off')\n\n\ndef augmentation_tester(augmentations, seed=42):\n    \"\"\"Box colors: violet - intact, orange - cut by the image border but kept,\n    red dashed - less than MIN_VISIBILITY of it left inside the image, would be dropped.\n    \"\"\"\n    rng = np.random.default_rng(seed)\n    sources = sorted(per_image['source'].unique())\n    fig, axes = plt.subplots(nrows=len(augmentations), ncols=2, figsize=(10, 5.3 * len(augmentations)), squeeze=False)\n\n    with torch.random.fork_rng(devices=[]):\n        torch.manual_seed(seed)\n        for row, (name, transform) in enumerate(augmentations.items()):\n            source = sources[row % len(sources)]\n            image_id = rng.choice(per_image.index[per_image['source'] == source])\n            boxes_before = bounding_boxes[rows_indices_by_image[image_id]]\n            colors_before = ['violet'] * len(boxes_before)\n            image = Image.open(f'{INPUT_DIR}/train/{image_id}.jpg')\n\n            augmented, boxes_after, visibility = augment(transform, image, boxes_before)\n\n            alive_boxes = ((boxes_after[:, 2] - boxes_after[:, 0]) >= 1) & (\n                    (boxes_after[:, 3] - boxes_after[:, 1]) >= 1)  # Note: Sanitizes bounding boxes\n            cut_boxes = alive_boxes & (visibility < 0.999)\n            dropped_boxes = cut_boxes & (visibility < MIN_VISIBILITY)\n            colors_mask_after = np.where(dropped_boxes, 'red', np.where(cut_boxes, 'orange', 'violet'))\n\n            geometric_area_mean = lambda boxes: np.median(np.sqrt(box_area(boxes))) if len(boxes) else float('nan')\n            max_dimension = max(image.width, image.height, augmented.width,\n                                augmented.height)  # Note: for cropped images it would not be IMAGE_SIZE\n\n            original_image_title_first_row = f'original  |  {source}  |  {image_id}'\n            original_image_title_second_row = f'{image.width}x{image.height}  |  {len(boxes_before)} boxes, median {geometric_area_mean(boxes_before):.0f}px'\n            draw_panel(ax=axes[row, 0], image=image, boxes=boxes_before, colors=colors_before,\n                       max_dimension=max_dimension,\n                       title=f'{original_image_title_first_row}\\n{original_image_title_second_row}')\n\n            kept = alive_boxes & ~dropped_boxes\n            augmented_image_title_first_row = f'{name}'\n            median_info = f'median {geometric_area_mean(boxes_after[kept]):.0f}px'\n            crop_info = f'  |  cut {cut_boxes.sum()}, dropped {dropped_boxes.sum()}' if cut_boxes.any() else ''\n            augmented_image_title_second_row = f'{augmented.width}x{augmented.height}  |  {kept.sum()} boxes, {median_info}{crop_info}'\n\n            draw_panel(ax=axes[row, 1], image=augmented, boxes=boxes_after[alive_boxes],\n                       colors=colors_mask_after[alive_boxes], max_dimension=max_dimension,\n                       title=f'{augmented_image_title_first_row}\\n{augmented_image_title_second_row}')\n\n    plt.tight_layout()\n    plt.show()\n\n\naugmentation_tester(TEST_AUGMENTATIONS)\n"},{"metadata":{},"cell_type":"markdown","source":"# Dataset preparation"},{"metadata":{},"cell_type":"markdown","source":"## New training dataset creation\nbased on the data analysis above"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:04:57.900488700Z","start_time":"2026-09-24T17:04:57.542023100Z"}},"cell_type":"code","source":"# Reinit as a sanity check\nall_image_ids = init_all_image_ids()\nrows_indices_by_image = init_rows_indices_by_image()\n\n\ndef init_is_duplicated_boxes():\n    \"\"\"\n    Runs on what survived the size filter and keeps the FIRST box of each duplicate group.\n    \"\"\"\n    is_duplicate = np.zeros(len(data_frame), dtype=bool)\n    for image_rows in rows_indices_by_image.values():\n        image_rows = image_rows[~too_big_mask.to_numpy()[image_rows]]\n        if len(image_rows) < 2 or data_frame['max_iou'].to_numpy()[image_rows].max() <= DUPLICATE_IOU:\n            continue\n        boxes = torch.from_numpy(bounding_boxes[image_rows])\n        iou = box_iou(boxes, boxes).fill_diagonal_(0).numpy()\n        kept_positions = []\n        for position in range(len(image_rows)):\n            if any(iou[position, kept] > DUPLICATE_IOU for kept in kept_positions):\n                is_duplicate[image_rows[position]] = True\n            else:\n                kept_positions.append(position)\n    return is_duplicate\n\n\ntoo_big_mask = data_frame['longest'] > TOO_BIG_LONGEST_SIDE\nis_duplicated_boxes = init_is_duplicated_boxes()\n\nremoval_reasons = (\n    pd.Series(data='', index=data_frame.index)\n    .case_when([\n        (too_big_mask, 'too big'),\n        (is_duplicated_boxes, 'duplicate')\n    ])\n)\nis_removed = removal_reasons != ''\nremoved_boxes = (data_frame.loc[is_removed, ['image_id', 'source', 'bbox', 'longest', 'max_iou']]\n                 .assign(reason=removal_reasons[is_removed]))\n\nkept_boxes = data_frame.loc[~is_removed, ORIGINAL_COLUMNS]\nimages_without_boxes = sorted(all_image_ids - set(kept_boxes['image_id']))\n\n# Note: because empty images have now a row, be cautious of you count boxes.\n#  Use `data_frame.groupby('image_id')['bbox'].count()` instead of `.size()`\nempty_image_rows = pd.DataFrame({\n    'image_id': images_without_boxes,\n    'width': ORIGINAL_IMAGE_SIZE,\n    'height': ORIGINAL_IMAGE_SIZE,\n    'bbox': np.nan,\n    'source': UNKNOWN_SOURCE,\n})\n\ncleaned_data_frame = (pd.concat([kept_boxes, empty_image_rows], ignore_index=True)\n                      .sort_values('image_id', kind='stable', ignore_index=True))\n\n# More sanity checks\nassert set(cleaned_data_frame['image_id']) == all_image_ids, 'csv and disk disagree about which images exist'\nassert cleaned_data_frame['bbox'].notna().sum() == len(data_frame) - len(removed_boxes), 'lost or gained a box'\nimage_ids_with_empty_row = set(cleaned_data_frame.loc[cleaned_data_frame['bbox'].isna(), 'image_id'])\nassert cleaned_data_frame['image_id'].isin(image_ids_with_empty_row).sum() == len(image_ids_with_empty_row), \\\n    'an image has both an empty row and real boxes'\n\ncleaned_data_frame.to_csv(INSPECT_ONLY_TRAIN_CSV_PATH, index=False)\n\nprint('=== Removed boxes ===')\nprint(removed_boxes.to_string(index=False))\n\nprint('\\n=== Cleaned dataset ===')\nprint(f'boxes     {len(data_frame)} -> {cleaned_data_frame[\"bbox\"].notna().sum()}  '\n      f'(-{(removal_reasons == \"too big\").sum()} too big, -{(removal_reasons == \"duplicate\").sum()} duplicate)')\nprint(f'images    {cleaned_data_frame[\"image_id\"].nunique()} (== files on disk), of them '\n      f'{len(images_without_boxes)} without boxes -> source \"{UNKNOWN_SOURCE}\", bbox NaN')\nprint(f'csv rows  {len(cleaned_data_frame)}')\nprint(f'saved     {INSPECT_ONLY_TRAIN_CSV_PATH}')\nprint('          (inspection only - the split below should use cleaned_data_frame, not this file)')\n","outputs":[],"execution_count":14},{"metadata":{},"cell_type":"markdown","source":"## Dataset split"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:04:58.514638200Z","start_time":"2026-09-24T17:04:57.913020800Z"}},"cell_type":"code","source":"N_FOLDS = 5\n# TODO: Change index, to see how validation metrics will change depending on different fold\n# Note: At last, try to learn on all folds, without validation to improve metrics\nVALIDATION_FOLD = 0\n\n# Note: Uses cleaned dataset with empty images from before\ncurrent_data_frame = cleaned_data_frame.copy()\nfolds = np.full(len(current_data_frame), -1)\n\nstratified_group_k_fold = StratifiedGroupKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)\n# Note: split() returns POSITIONAL indices, so they go into a numpy array.\n#  .loc would be label based and silently mismatch on any frame whose index is not a clean RangeIndex.\nsplit_data_frame_indices = stratified_group_k_fold.split(X=current_data_frame, y=current_data_frame['source'],\n                                                         groups=current_data_frame['image_id'])\n\nfor fold_num, (_, validation_indices) in enumerate(split_data_frame_indices):\n    folds[validation_indices] = fold_num\ncurrent_data_frame['fold'] = folds\n\nis_validation_mask = current_data_frame['fold'] == VALIDATION_FOLD\nvalidation_data_frame = current_data_frame[is_validation_mask].reset_index(drop=True)\ntrain_data_frame = current_data_frame[~is_validation_mask].reset_index(drop=True)\n\n#region Sanity checks & inspect output\nassert (folds >= 0).all(), 'some rows never landed in a fold'\nassert set(current_data_frame['image_id']) == all_image_ids, 'the split lost images'\nassert set(train_data_frame['image_id']).isdisjoint(validation_data_frame['image_id'])\nassert len(train_data_frame) + len(validation_data_frame) == len(current_data_frame)\n\n# Note: a row is no longer a box - an image without boxes is one row with bbox NaN\ncount_boxes = lambda frame: int(frame['bbox'].count())\ncount_empty_images = lambda frame: int(frame['bbox'].isna().sum())\n\nprint(f'validation fold {VALIDATION_FOLD} of {N_FOLDS}:  '\n      f'train {train_data_frame[\"image_id\"].nunique()} images / {count_boxes(train_data_frame)} boxes  |  '\n      f'validation {validation_data_frame[\"image_id\"].nunique()} images / {count_boxes(validation_data_frame)} boxes')\n# Note: an empty image scores 0 as soon as ANY box is predicted on it, so validation needs its share of them\nprint(f'images without boxes:  train {count_empty_images(train_data_frame)}  |  '\n      f'validation {count_empty_images(validation_data_frame)}')\n\nprint('\\n=== Images per source in each fold ===')\none_row_per_image = current_data_frame.drop_duplicates('image_id')\nprint(pd.crosstab(one_row_per_image['source'], one_row_per_image['fold'],\n                  margins=True, margins_name='total').to_string())\n\nprint('\\nboxes per fold: ', current_data_frame.groupby('fold')['bbox'].count().to_dict())\n#endregion","outputs":[],"execution_count":15},{"metadata":{},"cell_type":"markdown","source":"# Training loop"},{"metadata":{},"cell_type":"markdown","source":"## Training preparation"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:07:01.116999400Z","start_time":"2026-09-24T17:06:58.834294900Z"}},"cell_type":"code","source":"# TODO: Extract to high level constants?\nDETECTION_HEAD = 'faster_rcnn'  # 'faster_rcnn' | 'retinanet' | 'fcos'\nIMAGE_SIZE = ORIGINAL_IMAGE_SIZE  # Note: the detector's internal resize is a no-op at this size\nBATCH_SIZE = 4\nNUM_WORKERS = 4  # Note: Try 2 or 0?\nMAX_DETECTIONS_PER_IMAGE = 300\nCONFIDENCE_THRESHOLD = 0.55  # Note: where the METRIC is scored - swept on validation\n\n# Note: where the MODEL filters, before NMS. Kept low, bc three heads produce confidence on completely different scales\nMODEL_SCORE_THRESHOLD = 0.05\nNMS_THRESHOLD = 0.5  # TODO: needs tuning, bc 0.72% of ground-truth pairs already overlap above IoU 0.5\nIOU_THRESHOLDS = (0.50, 0.55, 0.60, 0.65, 0.70, 0.75)  # Note: as per competition rules\n\n# region Data\ntrain_transform = build_train_transform(image_size=IMAGE_SIZE, min_visibility=MIN_VISIBILITY)\n# Note: No augmentations and MIN_VISIBILITY=0 so that validation never loses a ground-truth box\nvalidation_transform = build_validation_transform(image_size=IMAGE_SIZE)\n\ntrain_dataset = ImageDataset(train_data_frame, image_dir=INPUT_DIR / 'train', transform=train_transform)\nvalidation_dataset = ImageDataset(validation_data_frame, image_dir=INPUT_DIR / 'train', transform=validation_transform)\n\ndata_loader_arguments = dict(\n    batch_size=BATCH_SIZE,\n    num_workers=NUM_WORKERS,\n    collate_fn=detection_collate,  # TODO: Learn more about collate_fn\n    pin_memory=True,\n    persistent_workers=NUM_WORKERS > 0\n)\n# Note: torchvision is doing preprocessing too, so there is no need to specify normalization step for this case\ntrain_loader = DataLoader(train_dataset, shuffle=True, drop_last=True, **data_loader_arguments)\n# Note: drop_last=False bc validation must score every image!\nvalidation_loader = DataLoader(validation_dataset, shuffle=False, drop_last=False, **data_loader_arguments)\n# endregion\n\n# region Model\nmodel = build_detector(DETECTION_HEAD, image_size=IMAGE_SIZE, num_classes=NUM_CLASSES,\n                       max_detections_per_image=MAX_DETECTIONS_PER_IMAGE,\n                       score_threshold=MODEL_SCORE_THRESHOLD, nms_threshold=NMS_THRESHOLD)\nmodel.to(DEVICE)\n# endregion\n\n# region Inspection & sanity checks\n# def print_model_info():\n#     print(describe_detector(model, DETECTION_HEAD))\n\n\ndef check_one_batch():\n    # Sanity check: one real batch must produce finite losses, including for the empty image\n    model.train()\n    images, targets = next(iter(train_loader))\n    with torch.autocast(DEVICE, dtype=torch.bfloat16):\n        losses = model([image.to(DEVICE) for image in images],\n                       [{key: value.to(DEVICE) if torch.is_tensor(value) else value for key, value in target.items()}\n                        for target in targets])\n    print(f'\\nlosses on one batch ({[len(target[\"boxes\"]) for target in targets]} boxes per image): '\n          f'{ {name: round(float(value), 4) for name, value in losses.items()} }')\n    assert all(torch.isfinite(value) for value in losses.values()), 'a loss came back non-finite'\n\n# print_model_info()\n# check_one_batch()\n# endregion","outputs":[],"execution_count":16},{"metadata":{},"cell_type":"markdown","source":"## Validation functions"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:07:01.138057900Z","start_time":"2026-09-24T17:07:01.118999300Z"}},"cell_type":"code","source":"@torch.inference_mode()\ndef predict(model, loader, device=DEVICE):\n    model.eval()\n    predictions, targets = [], []\n    for images, batch_targets in loader:\n        with torch.autocast(device, dtype=torch.bfloat16):\n            outputs = model([image.to(device, non_blocking=True) for image in images])\n\n        # Note: numpy does not support bfloat16.\n        predictions.extend({'boxes': output['boxes'].float().cpu(),\n                            'scores': output['scores'].float().cpu(),\n                            'labels': output['labels'].cpu()\n                            } for output in outputs)\n        targets.extend(batch_targets)\n    return predictions, targets\n\n\ndef evaluate(model, loader, confidence_threshold=CONFIDENCE_THRESHOLD):\n    predictions, targets = predict(model, loader)\n    return competition_metric(predictions, targets, confidence_threshold, thresholds=IOU_THRESHOLDS), predictions, targets","outputs":[],"execution_count":17},{"metadata":{},"cell_type":"markdown","source":"## Training loop"},{"metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:43:59.597090500Z","start_time":"2026-09-24T17:07:01.140058100Z"}},"cell_type":"code","source":"# region Training\nLEARNING_RATE = 1e-4  # TODO: Need more tuning\nWEIGHT_DECAY = 1e-4\nWARMUP_FRACTION = 0.05\nMAX_GRADIENT_NORM = 10.0\n\n\ndef build_lr_scheduler(optimizer, steps_per_epoch, epochs):\n    # Note: mostly reused from the first competition\n    total_steps = max(1, steps_per_epoch * epochs)\n    warmup_steps = max(1, int(WARMUP_FRACTION * total_steps))\n\n    def lr_factor(step):\n        if step < warmup_steps:\n            return (step + 1) / warmup_steps\n        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n        return 0.5 * (1 + math.cos(math.pi * progress))\n\n    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_factor)\n\n\ndef train_one_epoch(model, loader, optimizer, scheduler, device=DEVICE):\n    model.train()\n    totals, gradient_norms = {}, []\n    for images, targets in loader:\n        images = [image.to(device, non_blocking=True) for image in images]\n        targets = [{key: value.to(device) if torch.is_tensor(value) else value\n                    for key, value in target.items()} for target in targets]\n\n        with torch.autocast(device, dtype=torch.bfloat16):\n            losses = model(images, targets)\n        total_loss = sum(losses.values())\n\n        optimizer.zero_grad(set_to_none=True)\n        total_loss.backward()\n        # Note: clips gradients to partially mitigates exploding gradients\n        gradient_norms.append(float(torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRADIENT_NORM)))\n        optimizer.step()\n        scheduler.step()\n\n        for name, value in losses.items():\n            totals[name] = totals.get(name, 0.0) + float(value)\n        totals['loss'] = totals.get('loss', 0.0) + float(total_loss)\n\n    return {name: value / len(loader) for name, value in totals.items()}, float(np.mean(gradient_norms))\n\n\ndef format_losses(losses):\n    short = {'loss_classifier': 'cls', 'loss_box_reg': 'box', 'loss_objectness': 'obj',\n             'loss_rpn_box_reg': 'rpn', 'classification': 'cls', 'bbox_regression': 'box',\n             'bbox_ctrness': 'ctr'}\n    return ' '.join(f'{short.get(name, name.replace(\"loss_\", \"\"))} {value:.3f}'\n                    for name, value in losses.items() if name != 'loss')\n\n\ndef train(model, train_loader, validation_loader, epochs=EPOCHS,\n          learning_rate=LEARNING_RATE, checkpoint_path=None):\n    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=WEIGHT_DECAY)\n    scheduler = build_lr_scheduler(optimizer, len(train_loader), epochs)\n    metrics = {'train_loss': [], 'validation_metric': [], 'learning_rate': [],\n               'gradient_norm': [], 'seconds': []}\n    best_metric = -1.0\n\n    for epoch in range(epochs):\n        started = time.perf_counter()\n        current_learning_rate = optimizer.param_groups[0]['lr']\n        losses, gradient_norm = train_one_epoch(model, train_loader, optimizer, scheduler)\n        # TODO: Rewrite evaluate?\n        metric, _, _ = evaluate(model, validation_loader)\n        elapsed = time.perf_counter() - started\n\n        for name, value in losses.items():\n            metrics.setdefault('train_loss' if name == 'loss' else name, []).append(value)\n        metrics['validation_metric'].append(metric)\n        metrics['learning_rate'].append(current_learning_rate)\n        metrics['gradient_norm'].append(gradient_norm)\n        metrics['seconds'].append(elapsed)\n\n        improved = metric > best_metric\n        if improved:\n            best_metric = metric\n            if checkpoint_path is not None:\n                torch.save({'model': model.state_dict(),\n                            'epoch': epoch,\n                            'metric': metric,\n                            'metrics': metrics,\n                            'head': DETECTION_HEAD,\n                            'image_size': IMAGE_SIZE,\n                            'confidence_threshold': CONFIDENCE_THRESHOLD,\n                            'model_score_threshold': MODEL_SCORE_THRESHOLD,\n                            'max_detections_per_image': MAX_DETECTIONS_PER_IMAGE,\n                            'nms_threshold': NMS_THRESHOLD\n                            }, checkpoint_path)\n        print(f'ep {epoch:<3} loss {losses[\"loss\"]:.4f} ({format_losses(losses)})  '\n              f'grad {gradient_norm:6.2f}  lr {current_learning_rate:.1e}  '\n              f'val {metric:.4f}  {elapsed:5.0f}s' + ('  <- best' if improved else ''))\n    return metrics\n# endregion\n\n\nhistory = train(model, train_loader, validation_loader, epochs=EPOCHS, checkpoint_path=BEST_MODEL_PATH)","outputs":[],"execution_count":18},{"metadata":{},"cell_type":"markdown","source":"# Plot metrics"},{"cell_type":"code","metadata":{"ExecuteTime":{"end_time":"2026-09-24T17:44:35.577419200Z","start_time":"2026-09-24T17:43:59.607094900Z"}},"source":"def plot_training_curves(metrics):\n    fig, axes = plt.subplots(1, 3, figsize=(16, 4.5))\n    epochs = range(1, len(metrics['train_loss']) + 1)\n\n    axes[0].plot(epochs, metrics['train_loss'], 'o-', label='train')\n    axes[0].set_xlabel('epoch')\n    axes[0].set_ylabel('sum of the four losses')\n    axes[0].set_title('Training loss')\n    axes[0].legend()\n    axes[0].grid(alpha=.3)\n\n    # Note: not a fixed list - each head reports its own loss components\n    reserved = ('train_loss', 'validation_metric', 'learning_rate', 'gradient_norm', 'seconds')\n    drawn = [name for name in metrics if name not in reserved]\n    for name in drawn:\n        axes[1].plot(epochs, metrics[name], 'o-', label=name.replace('loss_', ''))\n    if drawn:\n        axes[1].legend()\n    else:\n        axes[1].text(0.5, 0.5, 'components not recorded in this run', transform=axes[1].transAxes,\n                     ha='center', va='center')\n    axes[1].set_xlabel('epoch')\n    axes[1].set_ylabel('loss')\n    axes[1].set_title('Loss components (cls/box = box head, obj/rpn = RPN)')\n    axes[1].grid(alpha=.3)\n\n    axes[2].plot(epochs, metrics['validation_metric'], 'o-', label='validation metric')\n\n    axes[2].set_xlabel('epoch')\n    axes[2].set_ylabel('competition metric')\n    axes[2].set_ylim(bottom=0)\n    axes[2].set_title(f'Competition metric (confidence {CONFIDENCE_THRESHOLD})')\n    axes[2].legend()\n    axes[2].grid(alpha=.3)\n\n    plt.tight_layout()\n    plt.show()\n\n\ndef plot_threshold_analysis(predictions, targets, candidates=np.arange(0.05, 1.0, 0.05),\n                            counted_iou=0.50):\n    by_confidence = [competition_metric(predictions, targets, candidate, IOU_THRESHOLDS)\n                     for candidate in candidates]\n    best_index = int(np.argmax(by_confidence))\n    best_confidence = float(candidates[best_index])\n\n    fig, axes = plt.subplots(1, 3, figsize=(16, 4.5))\n\n    axes[0].plot(candidates, by_confidence, 'o-')\n    axes[0].axvline(best_confidence, color='grey', ls='--', lw=1,\n                    label=f'best {by_confidence[best_index]:.4f} at {best_confidence:.2f}')\n    axes[0].axvline(CONFIDENCE_THRESHOLD, color='black', ls=':', lw=1,\n                    label=f'current {CONFIDENCE_THRESHOLD}')\n    axes[0].set_xlabel('confidence threshold')\n    axes[0].set_ylabel('competition metric')\n    axes[0].set_ylim(bottom=0)\n    axes[0].set_title('Metric vs confidence threshold')\n    axes[0].legend()\n    axes[0].grid(alpha=.3)\n\n    per_iou = [competition_metric(predictions, targets, best_confidence, (threshold,))\n               for threshold in IOU_THRESHOLDS]\n    axes[1].plot(IOU_THRESHOLDS, per_iou, 'o-')\n    axes[1].axhline(float(np.mean(per_iou)), color='grey', ls='--', lw=1,\n                    label=f'mean = the metric = {np.mean(per_iou):.4f}')\n    axes[1].set_xlabel('IoU threshold')\n    axes[1].set_ylabel('metric at this IoU alone')\n    # Note: from zero on purpose. Autoscaled, a 0.08 spread looks like a collapse\n    axes[1].set_ylim(0, max(max(per_iou), 0.01) * 1.2)\n    axes[1].set_title(f'Where the metric is lost (confidence {best_confidence:.2f})')\n    axes[1].legend()\n    axes[1].grid(alpha=.3)\n\n    counts = [match_counts(predictions, targets, candidate, counted_iou) for candidate in candidates]\n    for position, name in enumerate(('true positives', 'false positives', 'false negatives')):\n        axes[2].plot(candidates, [count[position] for count in counts], 'o-', label=name)\n    axes[2].set_xlabel('confidence threshold')\n    axes[2].set_ylabel(f'boxes over {len(targets)} images')\n    axes[2].set_title(f'What the metric is made of (IoU {counted_iou:.2f})')\n    axes[2].legend()\n    axes[2].grid(alpha=.3)\n\n    plt.tight_layout()\n    plt.show()\n    return best_confidence\n\n\ndef plot_per_source(predictions, targets, confidence_threshold):\n    scores = images_average_precision(predictions, targets, confidence_threshold, IOU_THRESHOLDS)\n    by_source = {}\n    for target, score in zip(targets, scores):\n        by_source.setdefault(target['source'], []).append(score)\n\n    sources = sorted(by_source, key=lambda name: float(np.mean(by_source[name])))\n    means = [float(np.mean(by_source[source])) for source in sources]\n    overall = float(np.mean(scores))\n\n    fig, ax = plt.subplots(figsize=(10, 4.5))\n    ax.barh(sources, means)\n    ax.axvline(overall, color='black', ls=':', lw=1, label=f'overall {overall:.4f}')\n    for position, source in enumerate(sources):\n        ax.text(means[position], position, f'  {means[position]:.3f} (n={len(by_source[source])})', va='center')\n    ax.set_xlabel('competition metric')\n    ax.set_title(f'Per source (confidence {confidence_threshold:.2f})')\n    ax.set_xlim(0, max(means) * 1.35)\n    ax.legend()\n    ax.grid(alpha=.3, axis='x')\n    plt.tight_layout()\n    plt.show()\n\n\n# Note: Plot metrics for the best epoch, not the last one\nmodel.load_state_dict(torch.load(BEST_MODEL_PATH, map_location=DEVICE, weights_only=True)['model'])\n\nmetric, predictions, targets = evaluate(model, validation_loader)\nprint(f'metric {metric:.4f} at confidence {CONFIDENCE_THRESHOLD}')\n\nplot_training_curves(history)\nbest_confidence = plot_threshold_analysis(predictions, targets)\nplot_per_source(predictions, targets, best_confidence)","outputs":[],"execution_count":19}]}