CoolFace
Datasetpublic

IOAI-official/IOAI2025

International Olympiad in Artificial Intelligence (IOAI 2025, Beijing, China) About IOAI 2025 The 2nd International Olympiad in Artificial Intelligence (IOAI 2025) took place in Beijing, China, from August 2 to 9, 2025, hosted by Beijing National Day School (BNDS) under the patronage of UNESCO. Contest Rules: Full rules encompassing the Individual, Team, and GAITE contests are available here. Syllabus: The official syllabus outlining the AI topics contestants… See the full description on the dataset page: https://huggingface.co/datasets/IOAI-official/IOAI2025.

sourceHugging Facecc-by-4.0updated 1y agoView on Hugging Face
1likes173downloads
Pixel_Solution.ipynb520 linesDownload Raw Back to Solution
1{2 "cells": [3  {4   "cell_type": "markdown",5   "id": "2051e891-2bcf-42a6-9a1e-f773baff8808",6   "metadata": {},7   "source": [8    "<img src=\"./figs/IOAI-Logo.png\" alt=\"IOAI Logo\" width=\"200\" height=\"auto\">\n",9    "\n",10    "[IOAI 2025 (Beijing, China), Individual Contest](https://ioai-official.org/china-2025)\n",11    "\n",12    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/IOAI-official/IOAI-2025/blob/main/Individual-Contest/Pixel/Solution/Pixel_Solution.ipynb)"13   ]14  },15  {16   "cell_type": "markdown",17   "id": "441c46a0-b20a-4dc1-9227-f659833a7d2f",18   "metadata": {},19   "source": [20    "# Pixel Efficiency: Reference Solution"21   ]22  },23  {24   "cell_type": "code",25   "execution_count": null,26   "id": "2df5ae63",27   "metadata": {},28   "outputs": [],29   "source": [30    "import os\n",31    "import numpy as np\n",32    "from PIL import Image\n",33    "from tqdm import tqdm\n",34    "import json\n",35    "import torch\n",36    "import torch.nn as nn\n",37    "from datasets import load_from_disk\n",38    "from transformers import CLIPProcessor, CLIPModel\n",39    "from typing import Optional\n",40    "from transformers.models.clip.modeling_clip import CLIPVisionTransformer, CLIPVisionConfig, BaseModelOutputWithPooling\n",41    "\n",42    "\n",43    "# Dataset configuration\n",44    "DATASET_PATH = os.environ.get(\"DATA_PATH\") + \"/test_dataset\"\n",45    "SPLIT = \"test\"\n",46    "\n",47    "# Model Configuration\n",48    "MODEL_PATH = \"./clip-vit-large-patch14\"\n",49    "DEVICE = \"cuda\"\n",50    "BACKGROUND_CLASS = \"other\"\n",51    "\n",52    "# Image and Masking Configuration\n",53    "HEIGHT = 224\n",54    "WIDTH = 224\n",55    "RETAIN_RATIO = 0.0625\n",56    "MEAN_COLOR = (0, 0, 0)\n",57    "STRIDE = 2\n",58    "TOP_K = 3\n",59    "\n",60    "# Load the dataset\n",61    "print(\"Loading dataset...\")\n",62    "dataset_whole = load_from_disk(DATASET_PATH)\n",63    "dataset = dataset_whole[SPLIT]\n",64    "\n",65    "print(f\"Dataset loaded successfully! Total samples: {len(dataset)}\")\n",66    "\n",67    "print(f\"Loading CLIP model and processor: {MODEL_PATH}...\")\n",68    "model = CLIPModel.from_pretrained(MODEL_PATH).to(DEVICE)\n",69    "processor = CLIPProcessor.from_pretrained(MODEL_PATH)\n",70    "print(\"Model and processor loaded successfully.\")\n",71    "\n",72    "\n",73    "def generate_all_rectangular_regions(image_size=224, patch_size=14, retain_ratio=RETAIN_RATIO, max_aspect_ratio=1.2, stride=1):\n",74    "    \"\"\"Generate rectangular regions with optimizations for speed\"\"\"\n",75    "    max_pixels = int(retain_ratio * image_size * image_size)\n",76    "    patches_per_side = image_size // patch_size\n",77    "    patch_area = patch_size * patch_size\n",78    "    target_patches = max_pixels // patch_area\n",79    "    \n",80    "    min_patches = max(1, target_patches - 1)\n",81    "    max_patches = target_patches + 1\n",82    "    \n",83    "    regions = []\n",84    "    region_to_patches = []\n",85    "    \n",86    "    # Pre-compute valid rectangle dimensions\n",87    "    valid_dims = []\n",88    "    for width_patches in range(1, patches_per_side + 1):\n",89    "        for height_patches in range(1, patches_per_side + 1):\n",90    "            total_patches = width_patches * height_patches\n",91    "            if min_patches <= total_patches <= max_patches:\n",92    "                aspect_ratio = max(width_patches, height_patches) / min(width_patches, height_patches)\n",93    "                if aspect_ratio <= max_aspect_ratio:\n",94    "                    valid_dims.append((width_patches, height_patches, total_patches))\n",95    "    \n",96    "    # Generate rectangles using stride for positions\n",97    "    for width_patches, height_patches, total_patches in valid_dims:\n",98    "        for top_patch in range(0, patches_per_side - height_patches + 1, stride):\n",99    "            for left_patch in range(0, patches_per_side - width_patches + 1, stride):\n",100    "                bottom_patch = top_patch + height_patches\n",101    "                right_patch = left_patch + width_patches\n",102    "                \n",103    "                pixel_coords = (\n",104    "                    top_patch * patch_size,\n",105    "                    left_patch * patch_size,\n",106    "                    bottom_patch * patch_size,\n",107    "                    right_patch * patch_size\n",108    "                )\n",109    "                regions.append(pixel_coords)\n",110    "                \n",111    "                covered_patches = []\n",112    "                for p_row in range(top_patch, bottom_patch):\n",113    "                    for p_col in range(left_patch, right_patch):\n",114    "                        patch_idx = p_row * patches_per_side + p_col\n",115    "                        covered_patches.append(patch_idx)\n",116    "                region_to_patches.append(covered_patches)\n",117    "    \n",118    "    return regions, region_to_patches\n",119    "\n",120    "\n",121    "class MaskCLIPVisionTransformer(CLIPVisionTransformer):\n",122    "    \"\"\"Modified CLIP Vision Transformer that supports mask tokens for all possible rectangular regions\"\"\"\n",123    "    \n",124    "    def __init__(self, config: CLIPVisionConfig, retain_ratio=RETAIN_RATIO):\n",125    "        super().__init__(config)\n",126    "        self.retain_ratio = retain_ratio\n",127    "        self.num_patches = (config.image_size // config.patch_size) ** 2\n",128    "        \n",129    "        self.regions, self.region_to_patches = generate_all_rectangular_regions(\n",130    "            image_size=config.image_size, \n",131    "            patch_size=config.patch_size, \n",132    "            retain_ratio=retain_ratio,\n",133    "            max_aspect_ratio=1.2,\n",134    "            stride=STRIDE\n",135    "        )\n",136    "        self.num_mask_tokens = len(self.regions)\n",137    "        \n",138    "        self.mask_tokens = nn.Parameter(torch.randn(1, self.num_mask_tokens, config.hidden_size))\n",139    "        \n",140    "    def create_mask_attention_matrix(self, batch_size):\n",141    "        \"\"\"Create attention mask matrix for all rectangular regions\"\"\"\n",142    "        N = self.num_patches\n",143    "        M = self.num_mask_tokens\n",144    "        total_tokens = N + 1 + M\n",145    "        \n",146    "        attention_mask = torch.zeros(total_tokens, total_tokens, dtype=torch.bool, device=self.mask_tokens.device)\n",147    "        \n",148    "        # Class token and image patches do NOT attend to mask tokens\n",149    "        attention_mask[:N+1, N+1:] = True\n",150    "        \n",151    "        # Each mask token attends to its specific image patches (not CLS)\n",152    "        attention_mask[N+1:, 1:N+1] = True\n",153    "        \n",154    "        # Then allow each mask token to attend to its assigned patches\n",155    "        for mask_idx in range(M):\n",156    "            covered_patches = self.region_to_patches[mask_idx]\n",157    "            for patch_idx in covered_patches:\n",158    "                token_pos = 1 + patch_idx\n",159    "                attention_mask[N + 1 + mask_idx, token_pos] = False\n",160    "        \n",161    "        # Mask tokens do NOT attend to each other\n",162    "        attention_mask[N+1:, N+1:] = True\n",163    "        # Allow self-attention for each mask token\n",164    "        for i in range(M):\n",165    "            attention_mask[N + 1 + i, N + 1 + i] = False\n",166    "        \n",167    "        return attention_mask\n",168    "    \n",169    "    def forward(\n",170    "        self,\n",171    "        pixel_values: Optional[torch.FloatTensor] = None,\n",172    "        output_attentions: Optional[bool] = None,\n",173    "        output_hidden_states: Optional[bool] = None,\n",174    "        interpolate_pos_encoding: Optional[bool] = False,\n",175    "        use_mask_tokens: bool = False,\n",176    "    ) -> BaseModelOutputWithPooling:\n",177    "        \"\"\"Forward pass with optional mask tokens\"\"\"\n",178    "        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n",179    "        output_hidden_states = (\n",180    "            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n",181    "        )\n",182    "\n",183    "        if pixel_values is None:\n",184    "            raise ValueError(\"You have to specify pixel_values\")\n",185    "\n",186    "        # Get embeddings (patches + class token)\n",187    "        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)\n",188    "        hidden_states = self.pre_layrnorm(hidden_states)\n",189    "        \n",190    "        if use_mask_tokens:\n",191    "            # Add mask tokens to the sequence\n",192    "            batch_size = hidden_states.shape[0]\n",193    "            \n",194    "            cls_token_embedding = hidden_states[:, 0:1, :]\n",195    "            mask_tokens_expanded = cls_token_embedding.expand(batch_size, self.num_mask_tokens, -1)\n",196    "            \n",197    "            if mask_tokens_expanded.device != hidden_states.device:\n",198    "                mask_tokens_expanded = mask_tokens_expanded.to(hidden_states.device)\n",199    "                \n",200    "            hidden_states = torch.cat([hidden_states, mask_tokens_expanded], dim=1)\n",201    "            \n",202    "            # Create custom attention mask\n",203    "            attention_mask = self.create_mask_attention_matrix(batch_size)\n",204    "            \n",205    "            seq_len = hidden_states.shape[1]\n",206    "            attention_mask_4d = attention_mask.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)\n",207    "            attention_mask_4d = attention_mask_4d.float()\n",208    "            attention_mask_4d = attention_mask_4d.masked_fill(attention_mask_4d == 1, float('-inf'))\n",209    "            attention_mask_4d = attention_mask_4d.masked_fill(attention_mask_4d == 0, 0.0)\n",210    "        else:\n",211    "            attention_mask_4d = None\n",212    "\n",213    "        # Process through encoder layers\n",214    "        encoder_outputs = self.encoder(\n",215    "            inputs_embeds=hidden_states,\n",216    "            attention_mask=attention_mask_4d,\n",217    "            causal_attention_mask=None,\n",218    "            output_attentions=output_attentions,\n",219    "            output_hidden_states=output_hidden_states,\n",220    "        )\n",221    "\n",222    "        last_hidden_state = encoder_outputs.last_hidden_state\n",223    "        \n",224    "        if use_mask_tokens:\n",225    "            # Extract different token types\n",226    "            class_token_output = last_hidden_state[:, 0]\n",227    "            mask_tokens_output = last_hidden_state[:, self.num_patches + 1:]\n",228    "            \n",229    "            # Apply post layer norm\n",230    "            pooled_output = self.post_layernorm(class_token_output)\n",231    "            mask_tokens_output = self.post_layernorm(mask_tokens_output)\n",232    "            \n",233    "            return {\n",234    "                'last_hidden_state': last_hidden_state,\n",235    "                'pooler_output': pooled_output,\n",236    "                'mask_tokens_output': mask_tokens_output,\n",237    "                'hidden_states': encoder_outputs.hidden_states,\n",238    "                'attentions': encoder_outputs.attentions,\n",239    "            }\n",240    "        else:\n",241    "            # Standard CLIP behavior\n",242    "            pooled_output = last_hidden_state[:, 0, :]\n",243    "            pooled_output = self.post_layernorm(pooled_output)\n",244    "\n",245    "            return BaseModelOutputWithPooling(\n",246    "                last_hidden_state=last_hidden_state,\n",247    "                pooler_output=pooled_output,\n",248    "                hidden_states=encoder_outputs.hidden_states,\n",249    "                attentions=encoder_outputs.attentions,\n",250    "            )\n",251    "\n",252    "\n",253    "def apply_mask_with_mean(image, mask, mean_rgb=MEAN_COLOR):\n",254    "    \"\"\"Apply arbitrary binary mask to image, replacing masked areas with mean values\"\"\"\n",255    "    img_array = np.array(image).copy()\n",256    "\n",257    "    if isinstance(mask, Image.Image):\n",258    "        mask_array = np.array(mask.convert('L')) > 127\n",259    "    else:\n",260    "        mask_array = mask > 0\n",261    "\n",262    "    mask_3d = np.stack([mask_array] * 3, axis=2)\n",263    "    mean_values = np.array([int(m * 255) for m in mean_rgb])\n",264    "    img_array = np.where(mask_3d, img_array, mean_values.reshape(1, 1, 3))\n",265    "\n",266    "    return Image.fromarray(img_array.astype(np.uint8))\n",267    "\n",268    "\n",269    "def compute_vision_features_once(model, image, mask_vision_model):\n",270    "    \"\"\"\n",271    "    Compute vision features once for efficient reuse across multiple mask selection functions.\n",272    "    This eliminates redundant forward passes.\n",273    "    \"\"\"\n",274    "    image_inputs = processor(images=image, return_tensors=\"pt\").to(DEVICE)\n",275    "    \n",276    "    with torch.no_grad():\n",277    "        vision_outputs = mask_vision_model(\n",278    "            pixel_values=image_inputs['pixel_values'],\n",279    "            use_mask_tokens=True\n",280    "        )\n",281    "        \n",282    "        full_image_features = vision_outputs['pooler_output']\n",283    "        if hasattr(model, 'visual_projection') and model.visual_projection is not None:\n",284    "            full_image_features = model.visual_projection(full_image_features)\n",285    "        \n",286    "        mask_tokens_features = vision_outputs['mask_tokens_output']\n",287    "        if hasattr(model, 'visual_projection') and model.visual_projection is not None:\n",288    "            batch_size, num_tokens, embed_dim = mask_tokens_features.shape\n",289    "            mask_tokens_features = mask_tokens_features.view(-1, embed_dim)\n",290    "            mask_tokens_features = model.visual_projection(mask_tokens_features)\n",291    "            mask_tokens_features = mask_tokens_features.view(batch_size, num_tokens, -1)\n",292    "        \n",293    "        full_image_features = full_image_features / full_image_features.norm(dim=-1, keepdim=True)\n",294    "        mask_tokens_features = mask_tokens_features / mask_tokens_features.norm(dim=-1, keepdim=True)\n",295    "    \n",296    "    return vision_outputs, full_image_features, mask_tokens_features\n",297    "\n",298    "\n",299    "def find_best_mask_region_calibrated(model, image, class_names, mask_vision_model, text_features, \n",300    "                                   vision_outputs=None, full_image_features=None, mask_tokens_features=None, \n",301    "                                   return_detailed=False):\n",302    "    \"\"\"Find the best mask region using MaskCLIP approach with calibration for black-pixel masking\"\"\"\n",303    "    num_mask_tokens = mask_vision_model.num_mask_tokens\n",304    "    \n",305    "    # Use pre-computed features if provided, otherwise compute them\n",306    "    if vision_outputs is None or full_image_features is None or mask_tokens_features is None:\n",307    "        image_inputs = processor(images=image, return_tensors=\"pt\").to(DEVICE)\n",308    "        \n",309    "        with torch.no_grad():\n",310    "            vision_outputs = mask_vision_model(\n",311    "                pixel_values=image_inputs['pixel_values'],\n",312    "                use_mask_tokens=True\n",313    "            )\n",314    "            \n",315    "            full_image_features = vision_outputs['pooler_output']\n",316    "            if hasattr(model, 'visual_projection') and model.visual_projection is not None:\n",317    "                full_image_features = model.visual_projection(full_image_features)\n",318    "            \n",319    "            mask_tokens_features = vision_outputs['mask_tokens_output']\n",320    "            if hasattr(model, 'visual_projection') and model.visual_projection is not None:\n",321    "                batch_size, num_tokens, embed_dim = mask_tokens_features.shape\n",322    "                mask_tokens_features = mask_tokens_features.view(-1, embed_dim)\n",323    "                mask_tokens_features = model.visual_projection(mask_tokens_features)\n",324    "                mask_tokens_features = mask_tokens_features.view(batch_size, num_tokens, -1)\n",325    "            \n",326    "            full_image_features = full_image_features / full_image_features.norm(dim=-1, keepdim=True)\n",327    "            mask_tokens_features = mask_tokens_features / mask_tokens_features.norm(dim=-1, keepdim=True)\n",328    "    \n",329    "    # Compute similarity between full image and text\n",330    "    full_image_similarities = torch.matmul(full_image_features, text_features.T)\n",331    "    full_image_prediction = torch.argmax(full_image_similarities, dim=-1)\n",332    "    predicted_class_idx = full_image_prediction.item()\n",333    "    \n",334    "    # Compute similarities for each mask token\n",335    "    mask_similarities = torch.matmul(mask_tokens_features.squeeze(0), text_features.T)\n",336    "    mask_predictions = torch.argmax(mask_similarities, dim=-1)\n",337    "    \n",338    "    # Get candidates that predict the same class as full image, sorted by confidence\n",339    "    matching_masks = (mask_predictions == predicted_class_idx)\n",340    "    \n",341    "    if matching_masks.any():\n",342    "        candidate_indices = torch.where(matching_masks)[0]\n",343    "        candidate_confidences = mask_similarities[candidate_indices, predicted_class_idx]\n",344    "        sorted_indices = torch.argsort(candidate_confidences, descending=True)\n",345    "        sorted_candidates = candidate_indices[sorted_indices]\n",346    "    else:\n",347    "        # If no exact matches, use all candidates sorted by confidence for predicted class\n",348    "        candidate_confidences = mask_similarities[:, predicted_class_idx]\n",349    "        sorted_candidates = torch.topk(candidate_confidences, len(candidate_confidences)).indices\n",350    "    \n",351    "    # OPTIMIZATION: If TOP_K=1, skip calibration and return best candidate directly\n",352    "    if TOP_K == 1:\n",353    "        return sorted_candidates[0].item()\n",354    "    \n",355    "    # CALIBRATION STEP: Test top K candidates, return immediately when one is correct\n",356    "    calibration_results = []\n",357    "    candidates_to_test = sorted_candidates[:TOP_K]\n",358    "    \n",359    "    for i, candidate_idx in enumerate(candidates_to_test):\n",360    "        candidate_idx_item = candidate_idx.item()\n",361    "        \n",362    "        # Create masked image for this candidate\n",363    "        coordinates = mask_idx_to_coordinates(candidate_idx_item, mask_vision_model)\n",364    "        mask = generate_mask_from_coordinates(image, coordinates)\n",365    "        masked_image = apply_mask_with_mean(image, mask)\n",366    "        \n",367    "        # Test with actual forward pass\n",368    "        with torch.no_grad():\n",369    "            masked_image_inputs = processor(images=masked_image, return_tensors=\"pt\").to(DEVICE)\n",370    "            masked_image_features = model.get_image_features(**masked_image_inputs)\n",371    "            masked_image_features = masked_image_features / masked_image_features.norm(dim=-1, keepdim=True)\n",372    "            \n",373    "            masked_similarities = torch.matmul(masked_image_features, text_features.T)\n",374    "            masked_prediction = torch.argmax(masked_similarities, dim=-1).item()\n",375    "            masked_confidence = masked_similarities[0, predicted_class_idx].item()\n",376    "        \n",377    "        # If this candidate predicts correctly, return it immediately (early exit optimization)\n",378    "        if masked_prediction == predicted_class_idx:\n",379    "            return candidate_idx_item\n",380    "        \n",381    "        # Store failed calibration result\n",382    "        calibration_results.append((candidate_idx_item, masked_confidence))\n",383    "    \n",384    "    # If we reach here, all TOP_K candidates failed calibration\n",385    "    # Fall back to the next best candidate from sorted list WITHOUT additional calibration\n",386    "    if len(sorted_candidates) > TOP_K:\n",387    "        return sorted_candidates[TOP_K].item()\n",388    "    else:\n",389    "        # If no more candidates available, return the best failed calibration result\n",390    "        if calibration_results:\n",391    "            return max(calibration_results, key=lambda x: x[1])[0]\n",392    "        else:\n",393    "            # Ultimate fallback: return the best mask token prediction\n",394    "            return sorted_candidates[0].item()\n",395    "\n",396    "\n",397    "def mask_idx_to_coordinates(mask_idx, mask_vision_model):\n",398    "    \"\"\"Convert mask token index to image coordinates using the pre-computed regions\"\"\"\n",399    "    if mask_idx >= len(mask_vision_model.regions):\n",400    "        raise ValueError(f\"mask_idx {mask_idx} is out of range. Only {len(mask_vision_model.regions)} regions available.\")\n",401    "    \n",402    "    top, left, bottom, right = mask_vision_model.regions[mask_idx]\n",403    "    return ((top, left), (bottom, right))\n",404    "\n",405    "\n",406    "def generate_mask_from_coordinates(image, coordinates):\n",407    "    \"\"\"Generate a binary mask from crop coordinates\"\"\"\n",408    "    H, W = 224, 224\n",409    "    mask = np.zeros((H, W), dtype=np.int8)\n",410    "    \n",411    "    (top, left), (bottom, right) = coordinates\n",412    "    mask[top:bottom, left:right] = 1\n",413    "    \n",414    "    return mask\n",415    "\n",416    "\n",417    "# Create the MaskCLIP model\n",418    "print(\"Creating MaskCLIP model...\")\n",419    "mask_vision_model = MaskCLIPVisionTransformer(model.vision_model.config, retain_ratio=RETAIN_RATIO)\n",420    "mask_vision_model.load_state_dict(model.vision_model.state_dict(), strict=False)\n",421    "mask_vision_model = mask_vision_model.to(DEVICE)\n",422    "mask_vision_model.eval()\n",423    "print(\"MaskCLIP model created successfully.\")\n",424    "\n",425    "dataset_eval = load_from_disk(DATASET_PATH)\n",426    "dataset_eval = dataset_eval[SPLIT]\n",427    "\n",428    "# Get class names from training dataset for consistent evaluation  \n",429    "train_dataset = load_from_disk(\"/bohr/train-yzfn/v1/train_dataset\")[\"train\"] # TODO: This part needs changing!\n",430    "class_names_eval = list(set([item['name'] for item in train_dataset])) + [BACKGROUND_CLASS]\n",431    "\n",432    "# Prepare text features once for efficiency\n",433    "print(\"Preparing text features...\")\n",434    "text_inputs_eval = processor(text=class_names_eval, return_tensors=\"pt\", padding=True).to(DEVICE)\n",435    "with torch.no_grad():\n",436    "    text_features_eval = model.get_text_features(**text_inputs_eval)\n",437    "    text_features_eval = text_features_eval / text_features_eval.norm(dim=-1, keepdim=True)\n",438    "print(\"Text features prepared.\")\n",439    "\n",440    "# Main evaluation loop\n",441    "masks = {}\n",442    "total_correct = 0\n",443    "total_processed = 0\n",444    "\n",445    "for item in tqdm(dataset_eval):\n",446    "    image = item['image']\n",447    "    total_processed += 1\n",448    "\n",449    "    try:\n",450    "        # Compute vision features once for efficiency (eliminates redundant forward passes)\n",451    "        vision_outputs, full_image_features, mask_tokens_features = compute_vision_features_once(\n",452    "            model, image, mask_vision_model\n",453    "        )\n",454    "        \n",455    "        # Get prediction from pre-computed features\n",456    "        full_image_similarities = torch.matmul(full_image_features, text_features_eval.T)\n",457    "        predicted_class_idx = torch.argmax(full_image_similarities, dim=-1).item()\n",458    "            \n",459    "        best_mask_idx = find_best_mask_region_calibrated(\n",460    "            model, image, class_names_eval, mask_vision_model, text_features_eval,\n",461    "            vision_outputs=vision_outputs, full_image_features=full_image_features, \n",462    "            mask_tokens_features=mask_tokens_features\n",463    "        )\n",464    "        \n",465    "        coordinates = mask_idx_to_coordinates(best_mask_idx, mask_vision_model)\n",466    "        \n",467    "        # Validate the mask\n",468    "        mask = generate_mask_from_coordinates(image, coordinates)\n",469    "        assert mask.shape == (224, 224), \"Mask should be 224x224\"\n",470    "        assert mask.sum() <= RETAIN_RATIO * 224 * 224, \"You should leave only 6.25% of pixels\"\n",471    "\n",472    "        \n",473    "        # Save the coordinates\n",474    "        idx = item['idx']\n",475    "        masks[idx] = coordinates\n",476    "        \n",477    "    except Exception as e:\n",478    "        print(f\"Error processing image {item['idx']}: {e}\")\n",479    "        # Fallback to a small center region if there's an error\n",480    "        if len(mask_vision_model.regions) > 0:\n",481    "            region_sizes = [(r[2]-r[0])*(r[3]-r[1]) for r in mask_vision_model.regions]\n",482    "            min_region_idx = region_sizes.index(min(region_sizes))\n",483    "            fallback_coords = mask_idx_to_coordinates(min_region_idx, mask_vision_model)\n",484    "        else:\n",485    "            fallback_coords = ((84, 84), (140, 140))\n",486    "        masks[item['idx']] = fallback_coords\n",487    "\n",488    "# Save as JSONL (one JSON object per line) - much safer than pickle\n",489    "with open('submission.jsonl', 'w') as f:\n",490    "    for idx, coordinates in masks.items():\n",491    "        json.dump({\"idx\": idx, \"coordinates\": coordinates}, f)\n",492    "        f.write('\\n')\n",493    "\n",494    "print(\"Masks saved to masks.jsonl\")\n"495   ]496  }497 ],498 "metadata": {499  "kernelspec": {500   "display_name": "Python 3 (ipykernel)",501   "language": "python",502   "name": "python3"503  },504  "language_info": {505   "codemirror_mode": {506    "name": "ipython",507    "version": 3508   },509   "file_extension": ".py",510   "mimetype": "text/x-python",511   "name": "python",512   "nbconvert_exporter": "python",513   "pygments_lexer": "ipython3",514   "version": "3.12.9"515  }516 },517 "nbformat": 4,518 "nbformat_minor": 5519}520