restructure_data.py

"""
Restructure dataset: raw images + JSON annotations.
No pre-extracted features. Contestants run DETR themselves.

Output:
  data_v2/
  ├── images/
  │   ├── train/     # training images
  │   ├── val/       # validation images  
  │   └── test/      # test images
  ├── train.json     # [{image_id, file_name, relationships: [{subject_box, object_box, predicate}]}]
  ├── val.json
  ├── test.json      # [{image_id, file_name}]  (no relationships)
  └── predicates.json
"""

import os
import json
import shutil
import random
import torch
import numpy as np

random.seed(2026)
np.random.seed(2026)

OLD_DATA = 'data'
NEW_DATA = 'data_v2'
IMG_SRC = os.path.join(OLD_DATA, 'images')

# Create structure
for split in ['train', 'val', 'test']:
    os.makedirs(os.path.join(NEW_DATA, 'images', split), exist_ok=True)

# Load old data
train_data = torch.load(f'{OLD_DATA}/train.pt', weights_only=False)
val_data = torch.load(f'{OLD_DATA}/val.pt', weights_only=False)
test_data = torch.load(f'{OLD_DATA}/test.pt', weights_only=False)
gt_data = torch.load(f'{OLD_DATA}/test_ground_truth.pt', weights_only=False)

with open(f'{OLD_DATA}/predicates.json', 'r') as f:
    pred_info = json.load(f)
predicates = pred_info['predicates']

# Build GT lookup for test
gt_lookup = {e['image_id']: e['triplets'] for e in gt_data}

def process_split(entries, split, has_triplets, gt_lookup=None):
    """Convert from .pt format to JSON + copy images."""
    annotations = []
    copied = 0
    skipped = 0
    
    for entry in entries:
        img_id = entry['image_id']
        src_path = os.path.join(IMG_SRC, f'{img_id}.jpg')
        dst_path = os.path.join(NEW_DATA, 'images', split, f'{img_id}.jpg')
        
        if not os.path.exists(src_path):
            skipped += 1
            continue
        
        # Copy image
        if not os.path.exists(dst_path):
            shutil.copy2(src_path, dst_path)
        copied += 1
        
        ann = {
            'image_id': img_id,
            'file_name': f'{img_id}.jpg',
        }
        
        if has_triplets:
            # Get triplets
            if gt_lookup and img_id in gt_lookup:
                triplets = gt_lookup[img_id]
            else:
                triplets = entry.get('triplets', [])
            
            # Convert triplets from detection indices to bounding boxes
            boxes = entry['det_boxes'].float().tolist()
            labels = entry['det_labels'].tolist()
            scores = entry['det_scores'].float().tolist()
            
            # Store detections for reference (ground truth objects)
            relationships = []
            for s_idx, o_idx, p_idx in triplets:
                if s_idx < len(boxes) and o_idx < len(boxes):
                    relationships.append({
                        'subject_box': boxes[s_idx],
                        'object_box': boxes[o_idx],
                        'subject_label': labels[s_idx],
                        'object_label': labels[o_idx],
                        'predicate': p_idx,
                        'predicate_name': predicates[p_idx],
                    })
            
            if len(relationships) > 0:
                ann['relationships'] = relationships
            else:
                skipped += 1
                copied -= 1
                continue
        
        annotations.append(ann)
    
    print(f"  {split}: {copied} images, {skipped} skipped")
    return annotations

print("Processing splits...")
train_ann = process_split(train_data, 'train', has_triplets=True)
val_ann = process_split(val_data, 'val', has_triplets=True)
test_ann = process_split(test_data, 'test', has_triplets=False)

# For test ground truth (organizer only)
test_ann_gt = process_split(test_data, 'test', has_triplets=True, gt_lookup=gt_lookup)

# Save JSON annotations
with open(os.path.join(NEW_DATA, 'train.json'), 'w') as f:
    json.dump(train_ann, f, indent=2)
print(f"  train.json: {len(train_ann)} entries")

with open(os.path.join(NEW_DATA, 'val.json'), 'w') as f:
    json.dump(val_ann, f, indent=2)
print(f"  val.json: {len(val_ann)} entries")

with open(os.path.join(NEW_DATA, 'test.json'), 'w') as f:
    json.dump(test_ann, f, indent=2)
print(f"  test.json: {len(test_ann)} entries (no relationships)")

# Ground truth for evaluation (organizer only, not given to contestants)
with open(os.path.join(NEW_DATA, 'test_ground_truth.json'), 'w') as f:
    json.dump(test_ann_gt, f, indent=2)
print(f"  test_ground_truth.json: {len(test_ann_gt)} entries")

# Predicates
with open(os.path.join(NEW_DATA, 'predicates.json'), 'w') as f:
    json.dump({
        'predicates': predicates,
        'n_predicates': len(predicates)
    }, f, indent=2)
print(f"  predicates.json: {len(predicates)} classes")

# Print sample
print(f"\nSample train entry:")
sample = train_ann[0]
print(f"  image_id: {sample['image_id']}")
print(f"  file_name: {sample['file_name']}")
print(f"  relationships ({len(sample['relationships'])}):")
for r in sample['relationships'][:3]:
    print(f"    {r['predicate_name']}: box {r['subject_box'][:2]}... -> box {r['object_box'][:2]}...")

# Total sizes
total_imgs = 0
for split in ['train', 'val', 'test']:
    d = os.path.join(NEW_DATA, 'images', split)
    n = len([f for f in os.listdir(d) if f.endswith('.jpg')])
    size = sum(os.path.getsize(os.path.join(d,f)) for f in os.listdir(d) if f.endswith('.jpg'))
    total_imgs += size
    print(f"  {split} images: {n} files, {size/1024/1024:.1f} MB")

print(f"\nTotal image size: {total_imgs/1024/1024:.1f} MB")
print("Done!")