File size: 4,259 Bytes
8222fd4 |
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 |
import json
import random
class FeatureCombiner():
def __init__(self, json_file, dataset=None) -> None:
with open(json_file, "r") as f:
data = json.load(f)
self.data = data
self.seg_data = [entry for entry in self.data if "annotations" in entry["image"]]
self.img_data = [entry for entry in self.data if "images" in entry["image"]]
self.img_data = [entry for entry in self.img_data if dataset in entry["dataset"]]
self.seg_data = [entry for entry in self.seg_data if dataset in entry["dataset"]]
self.dataset_field = dataset
if self.dataset_field is None:
self.feature_name = ["dataset", "location", "size", "trend_shape"]
else:
self.feature_name = ["location", "size", "trend_shape"]
def combine(self, is_seg_only=True, debug=False):
features_seg = []
features_img = []
# 选择seg_data的overview，这里假设每个entry中都有"overview"键
seg_overview_entry = random.choice(self.seg_data)
overview_seg = seg_overview_entry["overview"]
overview_img = self._find_img_entry_by_seg_entry(seg_overview_entry)["overview"]
features_seg.append(overview_seg)
features_img.append(overview_img)
if self.dataset_field != None:
features_seg.append(self.dataset_field)
features_img.append(self.dataset_field)
else:
raise "Not implement"
if not debug:
tmp_f_seg = []
tmp_f_img = []
# 根据is_seg_only选择数据集
data = self.seg_data if is_seg_only else self.data
for feature in self.feature_name:
# 随机选择特征，这里假设每个entry中都有这些feature_name中定义的键
entry_seg = random.choice(data)
entry_img = self._find_img_entry_by_seg_entry(entry_seg)
tmp_f_seg.append(entry_seg[feature])
tmp_f_img.append(entry_img[feature])
features_seg = features_seg + tmp_f_seg
features_img = features_img + tmp_f_img
caption_seg = ""
for sentence in features_seg:
# 整理sentence: 删除末尾的标点符号，去掉结尾多余的空格，然后在结尾添加分号，最后拼接到caption中
sentence = sentence.rstrip(';.!?') # 删除末尾的标点符号
sentence = sentence.strip() # 去掉结尾多余的空格
sentence += ";" # 在结尾添加分号
caption_seg += sentence + " " # 拼接到caption中，并添加一个空格以分隔句子
# 删除caption最后的分号和空格
caption_seg = caption_seg.rstrip('; ')
caption_img = ""
for sentence in features_img:
# 整理sentence: 删除末尾的标点符号，去掉结尾多余的空格，然后在结尾添加分号，最后拼接到caption中
sentence = sentence.rstrip(';.!?') # 删除末尾的标点符号
sentence = sentence.strip() # 去掉结尾多余的空格
sentence += ";" # 在结尾添加分号
caption_img += sentence + " " # 拼接到caption中，并添加一个空格以分隔句子
# 删除caption最后的分号和空格
caption_img = caption_img.rstrip('; ')
return caption_seg, caption_img
def _find_img_entry_by_seg_entry(self, seg_entry):
seg_path = seg_entry["image"]
file_name = seg_path.split('/')[-1].split(".")[0]
if "gt" in file_name:
file_name = file_name.replace("gt", "img")
for entry in self.img_data:
if file_name+"." in entry["image"]:
return entry
raise "No entry find"
if __name__ == "__main__":
json_file = r"/data/leiqin/dataset_curvilinear/dataset_curvilinear/crack/crack.json"
fc = FeatureCombiner(json_file,"Crack500 dataset")
for _ in range(9999):
caption1, caption2 = fc.combine()
print(caption1, caption2)
|