From 1b14779abc195ba721fa6a1fd6f9fb349c35f334 Mon Sep 17 00:00:00 2001
From: genitrix <mtrxcop@gmail.com>
Date: Tue, 18 Aug 2026 21:26:05 +0800
Subject: [PATCH] Upgrade OmniDocBench dataset to v1.5

---
 tests/test_omnidocbench.py                   | 39 ++++++++++++++++
 vlmeval/dataset/OmniDocBench/omnidocbench.py | 47 ++++++++++++++++----
 2 files changed, 78 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_omnidocbench.py

diff --git a/tests/test_omnidocbench.py b/tests/test_omnidocbench.py
new file mode 100644
index 000000000..0957a1500
--- /dev/null
+++ b/tests/test_omnidocbench.py
@@ -0,0 +1,39 @@
+import unittest
+
+from vlmeval.dataset.OmniDocBench.omnidocbench import OmniDocBench, parse_omnidocbench_annotation
+
+
+class OmniDocBenchTests(unittest.TestCase):
+
+    def test_v1_0_json_annotation(self):
+        value = '{"page_info": {"image_path": "page.png"}, "layout_dets": [], "extra": {}}'
+        parsed = parse_omnidocbench_annotation(value)
+        self.assertEqual(parsed['page_info']['image_path'], 'page.png')
+
+    def test_v1_5_python_literal_annotation(self):
+        value = "{'page_info': {'image_path': 'page.png'}, 'layout_dets': [], 'extra': {}}"
+        parsed = parse_omnidocbench_annotation(value)
+        self.assertEqual(parsed['page_info']['image_path'], 'page.png')
+
+    def test_non_dictionary_annotation_is_rejected(self):
+        with self.assertRaises(ValueError):
+            parse_omnidocbench_annotation("['not', 'an', 'annotation']")
+
+    def test_literal_fallback_cannot_execute_code(self):
+        with self.assertRaises(ValueError):
+            parse_omnidocbench_annotation("__import__('os').system('false')")
+
+    def test_default_dataset_is_pinned_v1_5(self):
+        self.assertIn(
+            '9702d4ba9a0d30dc5e76789707650c9c54cb0b3b',
+            OmniDocBench.DATASET_URL['OmniDocBench'],
+        )
+        self.assertEqual(
+            OmniDocBench.DATASET_MD5['OmniDocBench'],
+            '995f1af5b4e24ad0a6417cbff708b3fc',
+        )
+        self.assertIn('OmniDocBench_v1_0', OmniDocBench.supported_datasets())
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/vlmeval/dataset/OmniDocBench/omnidocbench.py b/vlmeval/dataset/OmniDocBench/omnidocbench.py
index a4a1fb188..c4701e614 100644
--- a/vlmeval/dataset/OmniDocBench/omnidocbench.py
+++ b/vlmeval/dataset/OmniDocBench/omnidocbench.py
@@ -1,3 +1,4 @@
+import ast
 import base64
 import copy
 import json
@@ -14,7 +15,26 @@
 from vlmeval.smp import dump, get_intermediate_file_path, load
 from ..image_base import ImageBaseDataset
 
-# from ..utils import get_intermediate_file_path, load, dump
+
+def parse_omnidocbench_annotation(value):
+    """Parse annotations from both official OmniDocBench TSV versions.
+
+    The v1.0 TSV stores ``answer`` as JSON, while the published v1.5 TSV
+    stores the same dictionary as a Python literal. ``ast.literal_eval`` is
+    intentionally used only as a fallback so no arbitrary code can execute.
+    """
+    if isinstance(value, dict):
+        return value
+    if not isinstance(value, str):
+        raise TypeError(f'Expected a string or dictionary, got {type(value).__name__}')
+
+    try:
+        annotation = json.loads(value)
+    except json.JSONDecodeError:
+        annotation = ast.literal_eval(value)
+    if not isinstance(annotation, dict):
+        raise ValueError(f'Expected an annotation dictionary, got {type(annotation).__name__}')
+    return annotation
 
 
 class OmniDocBench(ImageBaseDataset):
@@ -22,9 +42,20 @@ class OmniDocBench(ImageBaseDataset):
     MODALITY = 'IMAGE'
     TYPE = 'QA'
 
-    DATASET_URL = {'OmniDocBench':'https://huggingface.co/datasets/ouyanglinke/OmniDocBench_tsv/resolve/main/OmniDocBench.tsv'}
-    DATASET_MD5 = {'OmniDocBench': '0fa5ccf31e682e219cb9ca83da741a59'}
-
+    DATASET_URL = {
+        'OmniDocBench': (
+            'https://huggingface.co/datasets/ouyanglinke/OmniDocBench_tsv/resolve/'
+            '9702d4ba9a0d30dc5e76789707650c9c54cb0b3b/OmniDocBench.tsv'
+        ),
+        'OmniDocBench_v1_0': (
+            'https://huggingface.co/datasets/ouyanglinke/OmniDocBench_tsv/resolve/'
+            '5c4d32624f2f828ce73be900cb86d9e3fabaf760/OmniDocBench.tsv'
+        ),
+    }
+    DATASET_MD5 = {
+        'OmniDocBench': '995f1af5b4e24ad0a6417cbff708b3fc',
+        'OmniDocBench_v1_0': '0fa5ccf31e682e219cb9ca83da741a59',
+    }
 
     system_prompt = r'''You are an AI assistant specialized in converting PDF images to Markdown format. Please follow these instructions for the conversion:
 
@@ -102,10 +133,10 @@ def __init__(self,
         load_success,load_fail=0,0
         for i,ans in tqdm(enumerate(references),desc='Loading data'):
             try:
-                ans = json.loads(ans)
+                ans = parse_omnidocbench_annotation(ans)
                 load_success+=1
                 self.references.append(ans) #[{},{}]
-            except json.JSONDecodeError as e:
+            except (SyntaxError, TypeError, ValueError):
                 load_fail+=1
                 continue
         print(f'load_success:{load_success},load_fail:{load_fail}')
@@ -447,14 +478,14 @@ def load_data(self,eval_file,gt_file,pred_key,gt_key):
         load_success,load_fail=0,0
         for i,gt_sample in tqdm(enumerate(gt_samples),desc='Loading data'):
             try:
-                ans=json.loads(gt_sample)
+                ans=parse_omnidocbench_annotation(gt_sample)
                 for item in ans['layout_dets']:
                     if item['category_type']=="table":
                         item['pred']=predictions[i]
                         load_success+=1
                         preds.append(ans)
 
-            except json.JSONDecodeError as e:
+            except (SyntaxError, TypeError, ValueError):
                 load_fail+=1
                 continue
         print(f'load_table_success:{load_success},load_table_fail:{load_fail}')