TSGNet/dataset.py at master · jsgwer/TSGNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import os
import json
import torch
from torch.utils.data import Dataset
from PIL import Image

class ic_dataset_four_branch_fusion(Dataset):
    def __init__(self, txt_path, jsonl_path, img_path, tokenizer, transform=None, max_length=128):
        """
        Args:
            txt_path: 包含 img_name \t label 的文件路径
            jsonl_path: 统一的 JSONL 标签文件路径
            img_path: 图像文件夹路径
            tokenizer: 预训练的分词器
            transform: 图像预处理
            max_length: 文本最大长度
        """
        super().__init__()
        self.img_path = img_path
        self.transform = transform
        self.tokenizer = tokenizer
        self.max_length = max_length

        # 1. 解析基础数据列表 (图像名和分数)
        self.img_info_list = self.parse_txt(txt_path)

        # 2. 一次性加载 JSONL 数据并建立索引映射
        self.json_data_map = self.load_jsonl(jsonl_path)

    def parse_txt(self, txt_path):
        image_info_list = []
        with open(txt_path, 'r', encoding='utf-8') as f:
            for line in f:
                line_split = line.strip().split("\t")
                if len(line_split) >= 2:
                    image_info_list.append({
                        "name": line_split[0],
                        "label": float(line_split[1])
                    })
        return image_info_list

    def load_jsonl(self, jsonl_path):
        """将 JSONL 加载为字典，Key 为文件名"""
        data_map = {}
        with open(jsonl_path, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                # 建立文件名到分析内容的映射
                if "file_name" in item:
                    data_map[item["file_name"]] = item
        return data_map

    def _get_structured_text(self, file_name):
        """通过文件名获取并解析结构化文本"""
        # 如果 JSONL 中找不到该文件，使用预设的默认值
        json_item = self.json_data_map.get(file_name, {})
        analysis = json_item.get("analysis", {})

        # 1. Global (Summary)
        global_text = analysis.get("summary", "An image.")

        # 2. Numerical (Coverage)
        coverage_list = analysis.get("coverage", [])
        num_text = " . ".join(coverage_list) if coverage_list else "no detailed coverage."

        # 3. Spatial (Spatial)
        spatial_list = analysis.get("spatial", [])
        spatial_text = " [SEP] ".join(spatial_list) if spatial_list else "no spatial relations."

        # 4. Atypical (Atypical)
        atypical_list = analysis.get("atypical", [])
        atypical_text = " [SEP] ".join(atypical_list) if atypical_list else "typical scene"

        return global_text, num_text, spatial_text, atypical_text

    def _tokenize(self, text):
        return self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

    def __getitem__(self, index):
        item = self.img_info_list[index]
        img_name = item["name"]
        label = item["label"]

        # --- 1. 图像读取 ---
        img_full_path = os.path.join(self.img_path, img_name)
        img = Image.open(img_full_path).convert("RGB")
        if self.transform is not None:
            img = self.transform(img)

        # --- 2. 标签 ---
        label_tensor = torch.tensor(label, dtype=torch.float32)

        # --- 3. 从内存中的 Map 获取结构化文本 ---
        g_txt, n_txt, s_txt, a_txt = self._get_structured_text(img_name)

        # --- 4. 分支 Tokenization ---
        tokens_g = self._tokenize(g_txt)
        tokens_n = self._tokenize(n_txt)
        tokens_s = self._tokenize(s_txt)
        tokens_a = self._tokenize(a_txt)

        return {
            "img": img,
            "label": label_tensor,
            "g_ids": tokens_g['input_ids'].squeeze(0),
            "g_mask": tokens_g['attention_mask'].squeeze(0),
            "n_ids": tokens_n['input_ids'].squeeze(0),
            "n_mask": tokens_n['attention_mask'].squeeze(0),
            "s_ids": tokens_s['input_ids'].squeeze(0),
            "s_mask": tokens_s['attention_mask'].squeeze(0),
            "a_ids": tokens_a['input_ids'].squeeze(0),
            "a_mask": tokens_a['attention_mask'].squeeze(0)
        }

    def __len__(self):
        return len(self.img_info_list)