147 lines
4.8 KiB
Python
147 lines
4.8 KiB
Python
|
import os
|
|||
|
import pandas as pd
|
|||
|
import json
|
|||
|
|
|||
|
import requests
|
|||
|
|
|||
|
|
|||
|
|
|||
|
# #查看json文件
|
|||
|
# print_json_tree(data)
|
|||
|
|
|||
|
class image_caption_json():
|
|||
|
"""
|
|||
|
{
|
|||
|
"info": {...},
|
|||
|
"licenses": [...],
|
|||
|
"images": [...],
|
|||
|
"annotations": [...]
|
|||
|
}
|
|||
|
"""
|
|||
|
def __init__(self,json_path,data_number):
|
|||
|
self.json_path = json_path
|
|||
|
self.data_number = data_number
|
|||
|
with open(self.json_path, 'r') as f:
|
|||
|
self.data = json.load(f)
|
|||
|
|
|||
|
|
|||
|
# 树状输出json
|
|||
|
def print_json_tree(self, indent=1):
|
|||
|
for key, value in self.data.items():
|
|||
|
print(' ' * indent + str(key), end='\n')
|
|||
|
|
|||
|
# 如果值是列表,统计列表长度
|
|||
|
if isinstance(value, list):
|
|||
|
value_count = len(value)
|
|||
|
# 如果值是字典,统计字典的键数量
|
|||
|
elif isinstance(value, dict):
|
|||
|
value_count = len(value.keys())
|
|||
|
# 如果值是单个值,数量为1
|
|||
|
else:
|
|||
|
value_count = 1
|
|||
|
print(value_count)
|
|||
|
|
|||
|
|
|||
|
# if isinstance(value, dict):
|
|||
|
# print()
|
|||
|
# print_json_tree(value, indent + 1)
|
|||
|
# else:
|
|||
|
# print(': ' + str(value))
|
|||
|
|
|||
|
|
|||
|
|
|||
|
def download_image(self,url, save_dir, filename, timeout=10):
|
|||
|
# 1.创建保存路径
|
|||
|
if not os.path.exists(save_dir):
|
|||
|
os.makedirs(save_dir, exist_ok=True)
|
|||
|
try:
|
|||
|
response = requests.get(url, stream=True)
|
|||
|
response.raise_for_status() # 检查请求是否成功
|
|||
|
|
|||
|
img_path = f"{save_dir}/{filename}"
|
|||
|
with open(img_path, "wb") as f:
|
|||
|
for chunk in response.iter_content(1024):
|
|||
|
f.write(chunk)
|
|||
|
print(f"✅ Successfully saved {filename} to: {img_path} from: {url}\n")
|
|||
|
return img_path
|
|||
|
except Exception as e:
|
|||
|
print(f"❌ Failed to download from[{url}]: {e}\n")
|
|||
|
return None
|
|||
|
|
|||
|
def image_annotation(self,save_dir,csv_path):
|
|||
|
# if os.path.exists(save_dir):
|
|||
|
# print('coco_2014_caption目录已存在,跳过数据处理步骤')
|
|||
|
# return 0
|
|||
|
|
|||
|
|
|||
|
# 全局变量,记录最近成功的来源(0-"coco" 或 1-"flickr")
|
|||
|
LAST_SUCCESSFUL_SOURCE = 0
|
|||
|
# 初始化存储图片路径和描述的列表
|
|||
|
image_paths = []
|
|||
|
captions = []
|
|||
|
|
|||
|
# 获取前data_number个annotation
|
|||
|
for i, img_info in enumerate(self.data['images'][:self.data_number]):
|
|||
|
# 获取对应的caption和image
|
|||
|
img_id = img_info['id']
|
|||
|
filename = img_info['file_name']
|
|||
|
coco_url = img_info['coco_url']
|
|||
|
flickr_url = img_info['flickr_url']
|
|||
|
# 查找image对应的caption
|
|||
|
# # 匹配所有
|
|||
|
# match_annotation =[annotation['caption'] for annotation in self.data['annotations'] if annotation['image_id'] == img_id]
|
|||
|
# 只匹配第一个
|
|||
|
caption =next((annotation['caption'] for annotation in self.data['annotations'] if annotation['image_id'] == img_id),None)
|
|||
|
|
|||
|
print(f"{i+1}. 图片ID: {img_id}")
|
|||
|
print(f" 文件名: {filename}")
|
|||
|
print(f" Caption: {caption}")
|
|||
|
# print(f" coco_url: {coco_url}")
|
|||
|
|
|||
|
# 根据 url 下载图片
|
|||
|
if LAST_SUCCESSFUL_SOURCE:
|
|||
|
first_url = flickr_url
|
|||
|
second_url = coco_url
|
|||
|
else:
|
|||
|
first_url = coco_url
|
|||
|
second_url = flickr_url
|
|||
|
|
|||
|
image_path = self.download_image(first_url,save_dir,filename)
|
|||
|
if not image_path:
|
|||
|
image_path = self.download_image(second_url,save_dir,filename)
|
|||
|
if image_path:
|
|||
|
LAST_SUCCESSFUL_SOURCE =1-LAST_SUCCESSFUL_SOURCE
|
|||
|
else:
|
|||
|
print(f"❌❌ Failed to download ]\n")
|
|||
|
continue
|
|||
|
|
|||
|
# 将路径和描述添加到列表中
|
|||
|
image_paths.append(image_path)
|
|||
|
captions.append(caption)
|
|||
|
|
|||
|
# 将图片路径和描述保存为CSV文件
|
|||
|
df = pd.DataFrame({
|
|||
|
'image_path': image_paths,
|
|||
|
'caption': captions
|
|||
|
})
|
|||
|
# 将数据保存为CSV文件
|
|||
|
df.to_csv(csv_path, index=False)
|
|||
|
|
|||
|
print(f'数据处理完成')
|
|||
|
|
|||
|
|
|||
|
|
|||
|
|
|||
|
|
|||
|
|
|||
|
|
|||
|
if __name__ == '__main__':
|
|||
|
file_path = '/root/PMN_WS/coco_2014_caption/annotations/captions_train2014.json'
|
|||
|
MAX_DATA_NUMBER = 100
|
|||
|
|
|||
|
image_caption = image_caption_json(file_path,MAX_DATA_NUMBER)
|
|||
|
save_dir='/root/PMN_WS/qwen-test/coco_2014_image'
|
|||
|
csv_path = '/root/PMN_WS/qwen-test/coco-2024-dataset.csv'
|
|||
|
image_caption.image_annotation(save_dir,csv_path)
|
|||
|
# image_caption.print_json_tree()
|