| """Small CPU checks for release interfaces; no model assets or datasets required.""" |
| import argparse |
| from pathlib import Path |
| import tempfile |
| import gzip |
| import json |
| import unittest |
| from unittest.mock import patch |
| from types import SimpleNamespace |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from tools.config import parse_with_config, validate_paths |
| from tools.checkpoint import load_checkpoint, export_checkpoint |
| from eval.m4seg import greedy_mask_miou, image_size |
| from tools.outputs import image_mask_slots |
|
|
|
|
| class ConfigTests(unittest.TestCase): |
| def test_hub_inference_defaults_preserve_backbone_training_initialization(self): |
| from eval.tasks import parse_args as evaluation_args |
| from training.trainer import parse_args as training_args |
| from run import parse_args as inference_args |
| with patch('sys.argv', ['run.py']): |
| self.assertEqual(inference_args().merged_ckpt_path, 'PLAN-Lab/PRIMA') |
| self.assertEqual(evaluation_args([]).merged_ckpt_path, 'PLAN-Lab/PRIMA') |
| self.assertEqual(training_args([]).merged_ckpt_path, '') |
| self.assertEqual(evaluation_args(['--merged_ckpt_path', '']).merged_ckpt_path, '') |
|
|
| def test_cli_overrides_configuration(self): |
| with tempfile.TemporaryDirectory() as directory: |
| config = Path(directory) / 'options.yaml' |
| config.write_text('batch_size: 4\n') |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--batch_size', type=int, default=1) |
| args = parse_with_config(parser, ['--config', str(config), '--batch_size', '2']) |
| self.assertEqual(args.batch_size, 2) |
|
|
| def test_placeholders_fail_before_model_loading(self): |
| args = argparse.Namespace(version='/path/to/base_model') |
| with self.assertRaisesRegex(ValueError, 'replace the placeholder'): |
| validate_paths(args) |
|
|
| def test_complete_hub_model_does_not_require_separate_assets(self): |
| args = argparse.Namespace(merged_ckpt_path='PLAN-Lab/PRIMA', |
| version='/path/to/base_model', ckpt_path='/path/to/prima_weights') |
| validate_paths(args) |
| args.dataset_dir = '/path/to/datasets' |
| with self.assertRaisesRegex(ValueError, 'dataset_dir'): |
| validate_paths(args, dataset=True) |
|
|
| def test_complete_model_overlay_still_requires_an_existing_checkpoint(self): |
| args = argparse.Namespace(merged_ckpt_path='PLAN-Lab/PRIMA', |
| ckpt_path='/missing/prima_overlay.bin') |
| with self.assertRaisesRegex(ValueError, 'ckpt_path'): |
| validate_paths(args) |
|
|
| def test_routing_uses_actual_seg_id_without_leading_whitespace(self): |
| from tools.model_loader import setup_tokenizer_and_special_tokens, MODEL_DEFAULTS, INTERNAL_DEFAULTS |
|
|
| class Tokenizer: |
| unk_token = '<unk>' |
| unk_token_id = 0 |
|
|
| def convert_tokens_to_ids(self, token): |
| return 9 |
|
|
| def __call__(self, text, **kwargs): |
| return SimpleNamespace(input_ids=[5, 9, 7, 8]) |
|
|
| args = SimpleNamespace(**MODEL_DEFAULTS, **INTERNAL_DEFAULTS) |
| with patch('tools.model_loader.transformers.AutoTokenizer.from_pretrained', return_value=Tokenizer()): |
| setup_tokenizer_and_special_tokens(args) |
| self.assertEqual(args.seg_token_idx, 9) |
| self.assertTrue(all(pattern == [9, 7, 8] for pattern in args.seg_image_tokens)) |
|
|
|
|
| class PortableDataTests(unittest.TestCase): |
| def test_four_digit_objects_keep_mask_order_and_five_digit_parts(self): |
| from PIL import Image |
| from pycocotools import mask as mask_utils |
| from dataset.multi_image_datasets.M4seg_ds import M4SegDataset |
| from model.llava import conversation |
| with tempfile.TemporaryDirectory() as directory: |
| root = Path(directory) |
| for name in ('first.jpg', 'second.jpg'): |
| Image.new('RGB', (6, 4), 'white').save(root / name) |
| left = np.zeros((4, 6), dtype=np.uint8); left[:, :2] = 1 |
| right = np.zeros((4, 6), dtype=np.uint8); right[:, 4:] = 1 |
| def grounding(name, image, mask): |
| rle = mask_utils.encode(np.asfortranarray(mask)) |
| rle['counts'] = rle['counts'].decode('ascii') |
| return {'name': name, 'img': image, 'segmentation': rle} |
| record = {'pair_id': 'example', 'dataset_name': 'coco_objects', 'img_ids': [1, 2], |
| 'image_paths': ['first.jpg', 'second.jpg'], 'image_sizes': [[4, 6], [4, 6]], |
| 'q': 'Compare book_2149 and chair_201.', |
| 'a': 'The book_2149 and chair_201 have a leg_10102 nearby. Compare book_2149.', |
| 'groundings': {'img_1': [grounding('leg', 1, left)], |
| 'img_2': [grounding('book', 2, left), grounding('chair', 2, right)]}} |
| path = root / 'data.json'; path.write_text(json.dumps([record])) |
| dataset = M4SegDataset( |
| dataset_dir=root, |
| m4seg_annpath=path, mode='train') |
| with patch.object(conversation, 'default_conversation', conversation.conv_templates['llava_v1']): |
| item = dataset[0] |
| text = item[4][0] |
| self.assertIn('<p> book_149 </p> [SEG] (IMAGE2)', text) |
| self.assertIn('<p> leg </p> [SEG] (IMAGE1)', text) |
| self.assertEqual(text.count('<p> book_149 </p>'), 2) |
| self.assertNotIn('book_2149', text) |
| self.assertEqual(text.count('[SEG]'), 3) |
| self.assertTrue(np.array_equal(item[5][2][0].numpy(), left)) |
| self.assertTrue(np.array_equal(item[5][2][1].numpy(), right)) |
|
|
| |
| |
| record['a'] = 'The book_2149 repeats book_2149 before chair_201 and leg_10102.' |
| record['groundings'] = {'img_1': [], 'img_2': []} |
| import re |
| for match in re.finditer(r'book_2149|chair_201|leg_10102', record['a']): |
| name = match.group().split('_')[0] |
| image = 1 if name == 'leg' else 2 |
| entry = grounding(name, image, right if name == 'chair' else left) |
| entry.update(st_idx=match.start(), en_idx=match.end()) |
| record['groundings'][f'img_{image}'].append(entry) |
| path.write_text(json.dumps([record])) |
| dataset = M4SegDataset( |
| dataset_dir=root, |
| m4seg_annpath=path, mode='train') |
| with patch.object(conversation, 'default_conversation', conversation.conv_templates['llava_v1']): |
| item = dataset[0] |
| self.assertEqual(item[4][0].count('[SEG]'), 3) |
| self.assertEqual(item[4][0].count('<p> book_149 </p>'), 2) |
| self.assertTrue(np.array_equal(item[5][2][1].numpy(), right)) |
|
|
| def test_compressed_source_catalog_preserves_image_ids_without_mask_tables(self): |
| from dataset.utils.coco import COCO |
| from dataset.utils.utils import source_metadata_path |
| with tempfile.TemporaryDirectory() as directory: |
| root = Path(directory) |
| path = root / 'm4seg/metadata/source_catalogs/coco_objects/test.json.gz' |
| path.parent.mkdir(parents=True) |
| with gzip.open(path, 'wt') as handle: |
| json.dump({'images': [{'id': 17, 'file_name': '000000000017.jpg', |
| 'height': 3, 'width': 4}], |
| 'categories': [{'id': 2, 'name': 'object'}]}, handle) |
| resolved = source_metadata_path(root, 'coco_objects', 'val', root / 'missing.json') |
| api = COCO(resolved) |
| self.assertEqual(api.getImgIds(), [17]) |
| self.assertEqual(api.loadImgs([17])[0]['width'], 4) |
|
|
| def test_portable_gzip_annotations_need_no_source_annotation_files(self): |
| from dataset.multi_image_datasets.M4seg_ds import M4SegDataset |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / 'annotations.json.gz' |
| record = {'pair_id': 'sample', 'example_id': 'train-0000000', |
| 'dataset_name': 'coco_objects', 'img_ids': [1], |
| 'image_paths': ['coco_2017/train2017/example.jpg'], 'image_sizes': [[3, 4]], |
| 'q': 'Describe the image.', 'a': 'An image.', 'groundings': {'img_1': []}} |
| with gzip.open(path, 'wt') as handle: |
| json.dump([record], handle) |
| dataset = M4SegDataset( |
| dataset_dir=directory, |
| m4seg_annpath=path, mode='test') |
| self.assertTrue(dataset.portable_annotations) |
| self.assertFalse(dataset.data2api) |
| self.assertEqual(dataset.data_infos, [record]) |
|
|
| def test_portable_paths_reject_traversal(self): |
| from dataset.multi_image_datasets.M4seg_ds import portable_image_metadata |
| record = {'img_ids': [1], 'image_paths': ['../private.jpg'], 'image_sizes': [[3, 4]]} |
| with self.assertRaisesRegex(ValueError, 'safe relative'): |
| portable_image_metadata(record, '/path/to/images') |
|
|
| def test_unique_example_ids_preserve_repeated_source_pair_ids(self): |
| from eval.tasks import build_dataset |
| from tools.model_loader import MODEL_DEFAULTS, INTERNAL_DEFAULTS |
| args = SimpleNamespace(**MODEL_DEFAULTS, **INTERNAL_DEFAULTS, |
| dataset_dir='unused', m4seg_data='coco_objects', m4seg_annpath='unused') |
| records = [{'pair_id': 'repeated', 'example_id': 'train-0000000'}, |
| {'pair_id': 'repeated', 'example_id': 'train-0000001'}] |
| dataset = SimpleNamespace(data_infos=records) |
| with patch('eval.tasks.M4SegDataset', return_value=dataset): |
| self.assertIs(build_dataset(args, mode='train'), dataset) |
|
|
| def test_repeated_legacy_pair_ids_are_allowed_for_training_only(self): |
| from eval.tasks import build_dataset |
| from tools.model_loader import MODEL_DEFAULTS, INTERNAL_DEFAULTS |
| args = SimpleNamespace(**MODEL_DEFAULTS, **INTERNAL_DEFAULTS, |
| dataset_dir='unused', m4seg_data='coco_objects', m4seg_annpath='unused') |
| dataset = SimpleNamespace(data_infos=[{'pair_id': 'repeated'}, {'pair_id': 'repeated'}]) |
| with patch('eval.tasks.M4SegDataset', return_value=dataset): |
| self.assertIs(build_dataset(args, mode='train'), dataset) |
| with self.assertRaisesRegex(ValueError, 'unique'): |
| build_dataset(args, mode='test') |
|
|
|
|
| def small_model(): |
| outer = nn.Module() |
| outer.model = nn.Module() |
| core = outer.model |
| attention = nn.Module() |
| attention.lora_A = nn.Linear(2, 1, bias=False) |
| attention.lora_B = nn.Linear(1, 2, bias=False) |
| core.layers = nn.ModuleList([attention]) |
| core.grounding_encoder = nn.Module() |
| core.grounding_encoder.mask_decoder = nn.Linear(2, 2) |
| core.query_tokens = nn.Parameter(torch.ones(1, 2)) |
| core.residual_query_tokens = nn.Parameter(torch.ones(1, 2)) |
| core.residual_extractor = nn.Linear(2, 2) |
| core.mm_projector = nn.Linear(2, 2) |
| core.text_hidden_fcs = nn.Linear(2, 2) |
| core.visual_encoder = nn.Linear(2, 2) |
| core.visual_encoder.requires_grad_(False) |
| core.vision_tower = core.visual_encoder |
| return outer |
|
|
|
|
| class CheckpointTests(unittest.TestCase): |
| def test_export_omits_frozen_aliases_and_round_trips(self): |
| source = small_model() |
| destination = small_model() |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / 'weights.bin' |
| export_checkpoint(source, path) |
| state = torch.load(path, map_location='cpu') |
| self.assertFalse(any('visual_encoder' in key or 'vision_tower' in key for key in state)) |
| load_checkpoint(destination, path) |
| for key, value in state.items(): |
| self.assertTrue(torch.equal(destination.state_dict()[key], value), key) |
|
|
| def test_incomplete_checkpoint_is_rejected(self): |
| model = small_model() |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / 'weights.bin' |
| export_checkpoint(model, path) |
| state = torch.load(path, map_location='cpu') |
| del state['model.mm_projector.weight'] |
| torch.save(state, path) |
| with self.assertRaisesRegex(ValueError, 'Incomplete PRIMA checkpoint'): |
| load_checkpoint(model, path) |
|
|
| def test_peft_prefix_is_normalized(self): |
| model = small_model() |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / 'weights.bin' |
| export_checkpoint(model, path) |
| state = torch.load(path, map_location='cpu') |
| torch.save({'base_model.model.' + key: value for key, value in state.items()}, path) |
| self.assertGreater(load_checkpoint(small_model(), path)['loaded_tensors'], 0) |
|
|
|
|
| class MetricTests(unittest.TestCase): |
| def test_grounded_phrases_route_all_masks_and_normalize_labels(self): |
| from eval.m4seg import extract_labels |
| caption = ("<p> cabinet_17's drawer_2 </p> [SEG] (IMAGE1) [SEG] (IMAGE2) " |
| "<p> chair_12 </p> [SEG] (IMAGE2)") |
| self.assertEqual(extract_labels(caption, 3), |
| [['drawer'], ['drawer', 'chair'], []]) |
|
|
| def test_caption_cleaning_removes_grounding_markup(self): |
| from eval.m4seg import clean_caption, clean_reference_caption |
| self.assertEqual(clean_caption('The <p> chair_12 </p> [SEG] (IMAGE2) is red.'), |
| 'The chair is red.') |
| self.assertEqual(clean_reference_caption('The dining_chair_212 is red.'), |
| 'The dining chair is red.') |
|
|
| def test_direct_and_standalone_scoring_use_the_same_label_rules(self): |
| from eval.m4seg import evaluate, score_predictions |
|
|
| class Embedder: |
| revision = 'test' |
|
|
| def similarity(self, left, right): |
| return float(left == right) |
|
|
| grounding = {'name': 'chair_212', 'segmentation': [[0, 0, 4, 0, 4, 4, 0, 4]]} |
| records = [{'pair_id': 'example', 'dataset_name': 'source', 'img_ids': [1, 2], |
| 'a': 'A chair_101 and chair_201.', |
| 'groundings': {'img_1': [grounding], 'img_2': [grounding]}}] |
| with tempfile.TemporaryDirectory() as directory: |
| root = Path(directory) |
| prediction = root / 'example.pth' |
| torch.save({'output': ['<p> chair_101 </p> [SEG] (IMAGE1) [SEG] (IMAGE2)'], |
| 'pred_masks': [None, torch.ones(1, 4, 4), torch.ones(1, 4, 4)], |
| 'image_sizes': [[4, 4], [4, 4]]}, prediction) |
| direct = evaluate(root, {'example': prediction}, records, {'example'}, |
| embedder=Embedder()) |
| with patch('eval.m4seg.TextEmbedder', return_value=Embedder()): |
| standalone = score_predictions([root], records, with_text=True)['results'][0] |
| self.assertEqual(direct, standalone) |
| self.assertEqual(direct['miou'], 1.0) |
| self.assertEqual(direct['recall'], 1.0) |
| self.assertEqual(direct['semantic_similarity'], 1.0) |
| self.assertEqual(direct['iou_matched_semantic_similarity'], 1.0) |
|
|
| def test_polygon_sizes_come_from_native_metadata(self): |
| polygon = [{'segmentation': [[0, 0, 4, 0, 4, 3, 0, 3]]}] |
| self.assertEqual(image_size({'image_sizes': [[3, 4]]}, 0, polygon), (3, 4)) |
| self.assertIsNone(image_size({}, 0, polygon)) |
|
|
| def test_legacy_sizes_can_be_read_from_source_images(self): |
| from PIL import Image |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / 'image.png' |
| Image.new('RGB', (7, 5)).save(path) |
| self.assertEqual(image_size({'image_paths': [str(path)]}, 0, []), (5, 7)) |
|
|
| def test_matched_mask_denominator_is_preserved(self): |
| mask = np.eye(4, dtype=bool) |
| extra = np.zeros((4, 4), dtype=bool) |
| self.assertEqual(greedy_mask_miou([mask, extra], [mask]), 1.0) |
| self.assertEqual(greedy_mask_miou([], [mask]), 0.0) |
|
|
| def test_explicit_grounding_slots_keep_source_order(self): |
| first = torch.ones(1, 2, 2) |
| second = torch.zeros(1, 3, 3) |
| output = image_mask_slots([[None], [first], [second]], 2) |
| self.assertIs(output[0], first) |
| self.assertIs(output[1], second) |
| self.assertEqual(image_mask_slots(None, 2), [None, None]) |
|
|
| def test_evaluation_can_generate_and_score_in_one_command(self): |
| import json |
| from eval import tasks |
|
|
| records = [{'pair_id': 'example', 'dataset_name': 'source', 'img_ids': ['1'], |
| 'groundings': {'img_1': [{'name': 'item', |
| 'segmentation': [[0, 0, 4, 0, 4, 4, 0, 4]]}]}}] |
|
|
| class Dataset: |
| data_infos = records |
|
|
| def __len__(self): |
| return len(self.data_infos) |
|
|
| def generate(loader, model, tokenizer, args, output): |
| torch.save({'output': ['<p> item </p> [SEG] (IMAGE1)'], |
| 'pred_masks': [None, torch.ones(1, 4, 4)], |
| 'image_sizes': [[4, 4]]}, output / 'example.pth') |
| return {'samples': 1, 'written': 1, 'failures': {}} |
|
|
| with tempfile.TemporaryDirectory() as directory: |
| root = Path(directory) |
| annotation = root / 'annotations.json' |
| annotation.write_text(json.dumps(records)) |
| checkpoint = root / 'checkpoint.bin' |
| checkpoint.write_bytes(b'test') |
| args = tasks.parse_args(['--merged_ckpt_path', '', '--ckpt_path', str(checkpoint), '--m4seg_annpath', str(annotation), |
| '--output_save_path', str(root / 'outputs'), '--compute_metrics', |
| '--no-metrics_with_text', '--no-metrics_with_meteor']) |
| args.img_emb_len = 32 |
| with patch.object(tasks, 'validate_paths'), \ |
| patch.object(tasks, 'load_bundle', return_value=(object(), object(), object())), \ |
| patch.object(tasks, 'build_dataset', return_value=Dataset()), \ |
| patch.object(tasks, 'inference', side_effect=generate): |
| tasks.main(args) |
| report = json.loads((root / 'outputs/metrics.json').read_text()) |
| self.assertEqual(report['results'][0]['miou'], 1.0) |
| self.assertEqual(report['results'][0]['samples'], 1) |
|
|
| def test_generation_records_native_image_dimensions(self): |
| from eval import tasks |
|
|
| class Loader: |
| dataset = [0] |
|
|
| def __iter__(self): |
| yield {'sample_ids': ['example'], 'label_list': [(None,), (torch.zeros(3, 4),)], |
| 'image_paths': [['image.png']], 'conversation_list': ['question']} |
|
|
| model = SimpleNamespace(eval=lambda: None, |
| evaluate=lambda **kwargs: (torch.tensor([[1, 9, 2]]), |
| [[None], [torch.ones(1, 3, 4)]])) |
| tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=2, pad_token='<pad>', |
| batch_decode=lambda *args, **kwargs: ['answer']) |
| args = SimpleNamespace(local_rank=0, resume=False, max_new_tokens=16) |
| with tempfile.TemporaryDirectory() as directory: |
| output = Path(directory) |
| with patch.object(tasks, 'move_batch', side_effect=lambda batch, device: batch): |
| summary = tasks.inference(Loader(), model, tokenizer, args, output) |
| saved = torch.load(output / 'example.pth', map_location='cpu') |
| self.assertEqual(saved['image_sizes'], [(3, 4)]) |
| self.assertEqual(summary['failures'], {}) |
|
|
|
|
| class AttentionTests(unittest.TestCase): |
| def test_local_causal_mask_preserves_backbone_semantics(self): |
| from model.llava.model.language_model.prima_llama import PRIMAConfig, PRIMALlamaModel |
| from transformers.models.llama.modeling_llama import LlamaModel |
|
|
| config = PRIMAConfig(hidden_size=16, intermediate_size=32, num_attention_heads=4, |
| num_hidden_layers=1, vocab_size=32) |
| config._attn_implementation = 'eager' |
| model = PRIMALlamaModel(config) |
| inputs = torch.randn(1, 3, 16) |
| attention = torch.ones(1, 3) |
| positions = torch.arange(3) |
| expected = LlamaModel._update_causal_mask(model, attention, inputs, positions, None, False) |
| actual = model._update_causal_mask(attention, inputs, positions, None, False) |
| self.assertTrue(torch.equal(actual, expected)) |
| bf16 = model._update_causal_mask(attention, inputs.bfloat16(), positions, None, False) |
| self.assertEqual(bf16.dtype, torch.bfloat16) |
| self.assertEqual(bf16[0, 0, 0, 1].item(), torch.finfo(torch.bfloat16).min) |
|
|
|
|
| class TrainingTests(unittest.TestCase): |
| def test_wrapped_training_keeps_frozen_encoders_in_eval_mode(self): |
| from training.trainer import train_mode |
|
|
| class Model(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.core = nn.Module() |
| for name in ('visual_encoder', 'ln_vision', 'Qformer'): |
| setattr(self.core, name, nn.Linear(2, 2)) |
| self.core.grounding_encoder = nn.Module() |
| for name in ('image_encoder', 'prompt_encoder', 'mask_decoder'): |
| setattr(self.core.grounding_encoder, name, nn.Linear(2, 2)) |
|
|
| def get_model(self): |
| return self.core |
|
|
| wrapper = nn.Module() |
| wrapper.module = Model() |
| train_mode(wrapper) |
| self.assertFalse(wrapper.module.core.visual_encoder.training) |
| self.assertFalse(wrapper.module.core.Qformer.training) |
| self.assertFalse(wrapper.module.core.grounding_encoder.prompt_encoder.training) |
| self.assertTrue(wrapper.module.core.grounding_encoder.mask_decoder.training) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
