Pixel Efficiency

# Install required packages
!pip install numpy pillow tqdm torch transformers datasets accelerate

import os
import numpy as np
from PIL import Image
from tqdm import tqdm
import json
import torch
import torch.nn as nn
from datasets import load_dataset
from transformers import CLIPProcessor, CLIPModel
from typing import Optional
from transformers.models.clip.modeling_clip import CLIPVisionTransformer, CLIPVisionConfig, BaseModelOutputWithPooling


# Dataset configuration
DATASET_PATH = "IOAI-official/IOAI-2025-Pixel-test"
SPLIT = "test"

# Model Configuration
MODEL_PATH = "openai/clip-vit-large-patch14"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BACKGROUND_CLASS = "other"

# Image and Masking Configuration
HEIGHT = 224
WIDTH = 224
RETAIN_RATIO = 0.0625
MEAN_COLOR = (0, 0, 0)
STRIDE = 2
TOP_K = 3

# Load the dataset
print("Loading dataset...")
dataset = load_dataset(DATASET_PATH, split=SPLIT)

print(f"Dataset loaded successfully! Total samples: {len(dataset)}")

print(f"Loading CLIP model and processor: {MODEL_PATH}...")
model = CLIPModel.from_pretrained(MODEL_PATH).to(DEVICE)
processor = CLIPProcessor.from_pretrained(MODEL_PATH)
print("Model and processor loaded successfully.")


def generate_all_rectangular_regions(image_size=224, patch_size=14, retain_ratio=RETAIN_RATIO, max_aspect_ratio=1.2, stride=1):
    """Generate rectangular regions with optimizations for speed"""
    max_pixels = int(retain_ratio * image_size * image_size)
    patches_per_side = image_size // patch_size
    patch_area = patch_size * patch_size
    target_patches = max_pixels // patch_area
    
    min_patches = max(1, target_patches - 1)
    max_patches = target_patches + 1
    
    regions = []
    region_to_patches = []
    
    # Pre-compute valid rectangle dimensions
    valid_dims = []
    for width_patches in range(1, patches_per_side + 1):
        for height_patches in range(1, patches_per_side + 1):
            total_patches = width_patches * height_patches
            if min_patches <= total_patches <= max_patches:
                aspect_ratio = max(width_patches, height_patches) / min(width_patches, height_patches)
                if aspect_ratio <= max_aspect_ratio:
                    valid_dims.append((width_patches, height_patches, total_patches))
    
    # Generate rectangles using stride for positions
    for width_patches, height_patches, total_patches in valid_dims:
        for top_patch in range(0, patches_per_side - height_patches + 1, stride):
            for left_patch in range(0, patches_per_side - width_patches + 1, stride):
                bottom_patch = top_patch + height_patches
                right_patch = left_patch + width_patches
                
                pixel_coords = (
                    top_patch * patch_size,
                    left_patch * patch_size,
                    bottom_patch * patch_size,
                    right_patch * patch_size
                )
                regions.append(pixel_coords)
                
                covered_patches = []
                for p_row in range(top_patch, bottom_patch):
                    for p_col in range(left_patch, right_patch):
                        patch_idx = p_row * patches_per_side + p_col
                        covered_patches.append(patch_idx)
                region_to_patches.append(covered_patches)
    
    return regions, region_to_patches


class MaskCLIPVisionTransformer(CLIPVisionTransformer):
    """Modified CLIP Vision Transformer that supports mask tokens for all possible rectangular regions"""
    
    def __init__(self, config: CLIPVisionConfig, retain_ratio=RETAIN_RATIO):
        super().__init__(config)
        self.retain_ratio = retain_ratio
        self.num_patches = (config.image_size // config.patch_size) ** 2
        
        self.regions, self.region_to_patches = generate_all_rectangular_regions(
            image_size=config.image_size, 
            patch_size=config.patch_size, 
            retain_ratio=retain_ratio,
            max_aspect_ratio=1.2,
            stride=STRIDE
        )
        self.num_mask_tokens = len(self.regions)
        
        self.mask_tokens = nn.Parameter(torch.randn(1, self.num_mask_tokens, config.hidden_size))
        
    def create_mask_attention_matrix(self, batch_size):
        """Create attention mask matrix for all rectangular regions"""
        N = self.num_patches
        M = self.num_mask_tokens
        total_tokens = N + 1 + M
        
        attention_mask = torch.zeros(total_tokens, total_tokens, dtype=torch.bool, device=self.mask_tokens.device)
        
        # Class token and image patches do NOT attend to mask tokens
        attention_mask[:N+1, N+1:] = True
        
        # Each mask token attends to its specific image patches (not CLS)
        attention_mask[N+1:, 1:N+1] = True
        
        # Then allow each mask token to attend to its assigned patches
        for mask_idx in range(M):
            covered_patches = self.region_to_patches[mask_idx]
            for patch_idx in covered_patches:
                token_pos = 1 + patch_idx
                attention_mask[N + 1 + mask_idx, token_pos] = False
        
        # Mask tokens do NOT attend to each other
        attention_mask[N+1:, N+1:] = True
        # Allow self-attention for each mask token
        for i in range(M):
            attention_mask[N + 1 + i, N + 1 + i] = False
        
        return attention_mask
    
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        interpolate_pos_encoding: Optional[bool] = False,
        use_mask_tokens: bool = False,
    ) -> BaseModelOutputWithPooling:
        """Forward pass with optional mask tokens"""
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        # Get embeddings (patches + class token)
        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
        hidden_states = self.pre_layrnorm(hidden_states)
        
        if use_mask_tokens:
            # Add mask tokens to the sequence
            batch_size = hidden_states.shape[0]
            
            cls_token_embedding = hidden_states[:, 0:1, :]
            mask_tokens_expanded = cls_token_embedding.expand(batch_size, self.num_mask_tokens, -1)
            
            if mask_tokens_expanded.device != hidden_states.device:
                mask_tokens_expanded = mask_tokens_expanded.to(hidden_states.device)
                
            hidden_states = torch.cat([hidden_states, mask_tokens_expanded], dim=1)
            
            # Create custom attention mask
            attention_mask = self.create_mask_attention_matrix(batch_size)
            
            seq_len = hidden_states.shape[1]
            attention_mask_4d = attention_mask.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)
            attention_mask_4d = attention_mask_4d.float()
            attention_mask_4d = attention_mask_4d.masked_fill(attention_mask_4d == 1, float('-inf'))
            attention_mask_4d = attention_mask_4d.masked_fill(attention_mask_4d == 0, 0.0)
        else:
            attention_mask_4d = None

        # Process through encoder layers
        encoder_outputs = self.encoder(
            inputs_embeds=hidden_states,
            attention_mask=attention_mask_4d,
            causal_attention_mask=None,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        last_hidden_state = encoder_outputs.last_hidden_state
        
        if use_mask_tokens:
            # Extract different token types
            class_token_output = last_hidden_state[:, 0]
            mask_tokens_output = last_hidden_state[:, self.num_patches + 1:]
            
            # Apply post layer norm
            pooled_output = self.post_layernorm(class_token_output)
            mask_tokens_output = self.post_layernorm(mask_tokens_output)
            
            return {
                'last_hidden_state': last_hidden_state,
                'pooler_output': pooled_output,
                'mask_tokens_output': mask_tokens_output,
                'hidden_states': encoder_outputs.hidden_states,
                'attentions': encoder_outputs.attentions,
            }
        else:
            # Standard CLIP behavior
            pooled_output = last_hidden_state[:, 0, :]
            pooled_output = self.post_layernorm(pooled_output)

            return BaseModelOutputWithPooling(
                last_hidden_state=last_hidden_state,
                pooler_output=pooled_output,
                hidden_states=encoder_outputs.hidden_states,
                attentions=encoder_outputs.attentions,
            )


def apply_mask_with_mean(image, mask, mean_rgb=MEAN_COLOR):
    """Apply arbitrary binary mask to image, replacing masked areas with mean values"""
    img_array = np.array(image).copy()

    if isinstance(mask, Image.Image):
        mask_array = np.array(mask.convert('L')) > 127
    else:
        mask_array = mask > 0

    mask_3d = np.stack([mask_array] * 3, axis=2)
    mean_values = np.array([int(m * 255) for m in mean_rgb])
    img_array = np.where(mask_3d, img_array, mean_values.reshape(1, 1, 3))

    return Image.fromarray(img_array.astype(np.uint8))


def compute_vision_features_once(model, image, mask_vision_model):
    """
    Compute vision features once for efficient reuse across multiple mask selection functions.
    This eliminates redundant forward passes.
    """
    image_inputs = processor(images=image, return_tensors="pt").to(DEVICE)
    
    with torch.no_grad():
        vision_outputs = mask_vision_model(
            pixel_values=image_inputs['pixel_values'],
            use_mask_tokens=True
        )
        
        full_image_features = vision_outputs['pooler_output']
        if hasattr(model, 'visual_projection') and model.visual_projection is not None:
            full_image_features = model.visual_projection(full_image_features)
        
        mask_tokens_features = vision_outputs['mask_tokens_output']
        if hasattr(model, 'visual_projection') and model.visual_projection is not None:
            batch_size, num_tokens, embed_dim = mask_tokens_features.shape
            mask_tokens_features = mask_tokens_features.view(-1, embed_dim)
            mask_tokens_features = model.visual_projection(mask_tokens_features)
            mask_tokens_features = mask_tokens_features.view(batch_size, num_tokens, -1)
        
        full_image_features = full_image_features / full_image_features.norm(dim=-1, keepdim=True)
        mask_tokens_features = mask_tokens_features / mask_tokens_features.norm(dim=-1, keepdim=True)
    
    return vision_outputs, full_image_features, mask_tokens_features


def find_best_mask_region_calibrated(model, image, class_names, mask_vision_model, text_features, 
                                   vision_outputs=None, full_image_features=None, mask_tokens_features=None, 
                                   return_detailed=False):
    """Find the best mask region using MaskCLIP approach with calibration for black-pixel masking"""
    num_mask_tokens = mask_vision_model.num_mask_tokens
    
    # Use pre-computed features if provided, otherwise compute them
    if vision_outputs is None or full_image_features is None or mask_tokens_features is None:
        image_inputs = processor(images=image, return_tensors="pt").to(DEVICE)
        
        with torch.no_grad():
            vision_outputs = mask_vision_model(
                pixel_values=image_inputs['pixel_values'],
                use_mask_tokens=True
            )
            
            full_image_features = vision_outputs['pooler_output']
            if hasattr(model, 'visual_projection') and model.visual_projection is not None:
                full_image_features = model.visual_projection(full_image_features)
            
            mask_tokens_features = vision_outputs['mask_tokens_output']
            if hasattr(model, 'visual_projection') and model.visual_projection is not None:
                batch_size, num_tokens, embed_dim = mask_tokens_features.shape
                mask_tokens_features = mask_tokens_features.view(-1, embed_dim)
                mask_tokens_features = model.visual_projection(mask_tokens_features)
                mask_tokens_features = mask_tokens_features.view(batch_size, num_tokens, -1)
            
            full_image_features = full_image_features / full_image_features.norm(dim=-1, keepdim=True)
            mask_tokens_features = mask_tokens_features / mask_tokens_features.norm(dim=-1, keepdim=True)
    
    # Compute similarity between full image and text
    full_image_similarities = torch.matmul(full_image_features, text_features.T)
    full_image_prediction = torch.argmax(full_image_similarities, dim=-1)
    predicted_class_idx = full_image_prediction.item()
    
    # Compute similarities for each mask token
    mask_similarities = torch.matmul(mask_tokens_features.squeeze(0), text_features.T)
    mask_predictions = torch.argmax(mask_similarities, dim=-1)
    
    # Get candidates that predict the same class as full image, sorted by confidence
    matching_masks = (mask_predictions == predicted_class_idx)
    
    if matching_masks.any():
        candidate_indices = torch.where(matching_masks)[0]
        candidate_confidences = mask_similarities[candidate_indices, predicted_class_idx]
        sorted_indices = torch.argsort(candidate_confidences, descending=True)
        sorted_candidates = candidate_indices[sorted_indices]
    else:
        # If no exact matches, use all candidates sorted by confidence for predicted class
        candidate_confidences = mask_similarities[:, predicted_class_idx]
        sorted_candidates = torch.topk(candidate_confidences, len(candidate_confidences)).indices
    
    # OPTIMIZATION: If TOP_K=1, skip calibration and return best candidate directly
    if TOP_K == 1:
        return sorted_candidates[0].item()
    
    # CALIBRATION STEP: Test top K candidates, return immediately when one is correct
    calibration_results = []
    candidates_to_test = sorted_candidates[:TOP_K]
    
    for i, candidate_idx in enumerate(candidates_to_test):
        candidate_idx_item = candidate_idx.item()
        
        # Create masked image for this candidate
        coordinates = mask_idx_to_coordinates(candidate_idx_item, mask_vision_model)
        mask = generate_mask_from_coordinates(image, coordinates)
        masked_image = apply_mask_with_mean(image, mask)
        
        # Test with actual forward pass
        with torch.no_grad():
            masked_image_inputs = processor(images=masked_image, return_tensors="pt").to(DEVICE)
            masked_image_features = model.get_image_features(**masked_image_inputs)
            masked_image_features = masked_image_features / masked_image_features.norm(dim=-1, keepdim=True)
            
            masked_similarities = torch.matmul(masked_image_features, text_features.T)
            masked_prediction = torch.argmax(masked_similarities, dim=-1).item()
            masked_confidence = masked_similarities[0, predicted_class_idx].item()
        
        # If this candidate predicts correctly, return it immediately (early exit optimization)
        if masked_prediction == predicted_class_idx:
            return candidate_idx_item
        
        # Store failed calibration result
        calibration_results.append((candidate_idx_item, masked_confidence))
    
    # If we reach here, all TOP_K candidates failed calibration
    # Fall back to the next best candidate from sorted list WITHOUT additional calibration
    if len(sorted_candidates) > TOP_K:
        return sorted_candidates[TOP_K].item()
    else:
        # If no more candidates available, return the best failed calibration result
        if calibration_results:
            return max(calibration_results, key=lambda x: x[1])[0]
        else:
            # Ultimate fallback: return the best mask token prediction
            return sorted_candidates[0].item()


def mask_idx_to_coordinates(mask_idx, mask_vision_model):
    """Convert mask token index to image coordinates using the pre-computed regions"""
    if mask_idx >= len(mask_vision_model.regions):
        raise ValueError(f"mask_idx {mask_idx} is out of range. Only {len(mask_vision_model.regions)} regions available.")
    
    top, left, bottom, right = mask_vision_model.regions[mask_idx]
    return ((top, left), (bottom, right))


def generate_mask_from_coordinates(image, coordinates):
    """Generate a binary mask from crop coordinates"""
    H, W = 224, 224
    mask = np.zeros((H, W), dtype=np.int8)
    
    (top, left), (bottom, right) = coordinates
    mask[top:bottom, left:right] = 1
    
    return mask


# Create the MaskCLIP model
print("Creating MaskCLIP model...")
mask_vision_model = MaskCLIPVisionTransformer(model.vision_model.config, retain_ratio=RETAIN_RATIO)
mask_vision_model.load_state_dict(model.vision_model.state_dict(), strict=False)
mask_vision_model = mask_vision_model.to(DEVICE)
mask_vision_model.eval()
print("MaskCLIP model created successfully.")

# Get class names from training dataset for consistent evaluation  
train_dataset = load_dataset("IOAI-official/IOAI-2025-Pixel-train", split="train")
class_names_eval = list(set([item['name'] for item in train_dataset])) + [BACKGROUND_CLASS]

# Prepare text features once for efficiency
print("Preparing text features...")
text_inputs_eval = processor(text=class_names_eval, return_tensors="pt", padding=True).to(DEVICE)
with torch.no_grad():
    text_features_eval = model.get_text_features(**text_inputs_eval)
    text_features_eval = text_features_eval / text_features_eval.norm(dim=-1, keepdim=True)
print("Text features prepared.")

# Main evaluation loop
masks = {}
total_correct = 0
total_processed = 0

for item in tqdm(dataset):
    image = item['image']
    total_processed += 1

    try:
        # Compute vision features once for efficiency (eliminates redundant forward passes)
        vision_outputs, full_image_features, mask_tokens_features = compute_vision_features_once(
            model, image, mask_vision_model
        )
        
        # Get prediction from pre-computed features
        full_image_similarities = torch.matmul(full_image_features, text_features_eval.T)
        predicted_class_idx = torch.argmax(full_image_similarities, dim=-1).item()
            
        best_mask_idx = find_best_mask_region_calibrated(
            model, image, class_names_eval, mask_vision_model, text_features_eval,
            vision_outputs=vision_outputs, full_image_features=full_image_features, 
            mask_tokens_features=mask_tokens_features
        )
        
        coordinates = mask_idx_to_coordinates(best_mask_idx, mask_vision_model)
        
        # Validate the mask
        mask = generate_mask_from_coordinates(image, coordinates)
        assert mask.shape == (224, 224), "Mask should be 224x224"
        assert mask.sum() <= RETAIN_RATIO * 224 * 224, "You should leave only 6.25% of pixels"

        
        # Save the coordinates
        idx = item['idx']
        masks[idx] = coordinates
        
    except Exception as e:
        print(f"Error processing image {item['idx']}: {e}")
        # Fallback to a small center region if there's an error
        if len(mask_vision_model.regions) > 0:
            region_sizes = [(r[2]-r[0])*(r[3]-r[1]) for r in mask_vision_model.regions]
            min_region_idx = region_sizes.index(min(region_sizes))
            fallback_coords = mask_idx_to_coordinates(min_region_idx, mask_vision_model)
        else:
            fallback_coords = ((84, 84), (140, 140))
        masks[item['idx']] = fallback_coords

# Save as JSONL (one JSON object per line) - much safer than pickle
with open('submission.jsonl', 'w') as f:
    for idx, coordinates in masks.items():
        json.dump({"idx": idx, "coordinates": coordinates}, f)
        f.write('\n')

print("Masks saved to masks.jsonl")