|
|
import copy |
|
|
import os |
|
|
import json |
|
|
import random |
|
|
from PIL import Image, ImageDraw |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
def find_nearest(anno_list, idx): |
|
|
while anno_list[idx] is None: |
|
|
idx -= 1 |
|
|
return idx |
|
|
|
|
|
|
|
|
def parse_anno(path): |
|
|
with open(path, 'r') as f: |
|
|
datas = f.readlines() |
|
|
|
|
|
ret = [] |
|
|
for data in datas: |
|
|
bbox = data.replace("\n", "").split() |
|
|
bbox = [float(_item.strip()) for _item in bbox] |
|
|
ret.append(bbox) |
|
|
return ret |
|
|
|
|
|
|
|
|
if not os.path.exists('./achieved'): |
|
|
os.mkdir('./achieved') |
|
|
if not os.path.exists('./achieved/images/'): |
|
|
os.mkdir('./achieved/images') |
|
|
|
|
|
save_image_path = './achieved/images' |
|
|
save_json_path = './achieved/anno.json' |
|
|
|
|
|
final_json_data = { |
|
|
"task": "video object tracking, under water video", |
|
|
"data_source": "UTB180", |
|
|
"type": "comprehension", |
|
|
"modality": { |
|
|
"in": ["image", "text"], |
|
|
"out": ["text"] |
|
|
}, |
|
|
"version": 1.0, |
|
|
} |
|
|
|
|
|
src_frames_folder = 'UTB180/UTB180/' |
|
|
src_first_frame_folder = 'thumbnails/thumbnails/UTB180_1st_frames/' |
|
|
|
|
|
_PER_NUMBER=50 |
|
|
_SAMPLE_FRAMES=30 |
|
|
|
|
|
with open("./under_water_attr.txt", 'r') as f: |
|
|
datas = f.readlines() |
|
|
split_data_list = {'blue_water': [], "green_water": [], "yellow_water": [], "white_water": []} |
|
|
split_attr_list = {'blue_water': [], "green_water": [], "yellow_water": [], "white_water": []} |
|
|
for idx in range(1, 181): |
|
|
_str_idx = str(idx + 10000)[1:] |
|
|
idx_attr = datas[idx-1].lower() |
|
|
if 'blue' in idx_attr: |
|
|
split_attr_list['blue_water'].append(_str_idx) |
|
|
elif 'clear' in idx_attr or idx_attr == "": |
|
|
split_attr_list["white_water"].append(_str_idx) |
|
|
elif 'green' in idx_attr: |
|
|
split_attr_list["green_water"].append(_str_idx) |
|
|
elif 'brown' in idx_attr or 'brown' in idx_attr: |
|
|
split_attr_list["yellow_water"].append(_str_idx) |
|
|
|
|
|
_id = 10000 |
|
|
for instance_name in os.listdir(src_frames_folder): |
|
|
if '.json' in instance_name or '.xlsx' in instance_name: |
|
|
continue |
|
|
_split = None |
|
|
_sub_nums = None |
|
|
for _key_str in split_attr_list['blue_water']: |
|
|
if _key_str in instance_name: |
|
|
_split = 'blue_water' |
|
|
_sub_nums = _PER_NUMBER // len(split_attr_list['blue_water']) + 1 |
|
|
break |
|
|
if _split is None: |
|
|
for _key_str in split_attr_list['green_water']: |
|
|
if _key_str in instance_name: |
|
|
_split = 'green_water' |
|
|
_sub_nums = _PER_NUMBER // len(split_attr_list['green_water']) + 1 |
|
|
break |
|
|
if _split is None: |
|
|
for _key_str in split_attr_list['white_water']: |
|
|
print(_key_str, ' ', instance_name) |
|
|
if _key_str in instance_name: |
|
|
_split = 'white_water' |
|
|
_sub_nums = _PER_NUMBER // len(split_attr_list['white_water']) + 1 |
|
|
break |
|
|
if _split is None: |
|
|
for _key_str in split_attr_list['yellow_water']: |
|
|
if _key_str in instance_name: |
|
|
_split = 'yellow_water' |
|
|
_sub_nums = _PER_NUMBER // len(split_attr_list['yellow_water']) + 1 |
|
|
break |
|
|
if _split is None: |
|
|
continue |
|
|
|
|
|
if len(split_data_list[_split]) >= _PER_NUMBER: |
|
|
continue |
|
|
|
|
|
anno_file_path = os.path.join(src_frames_folder, instance_name, "groundtruth_rect.txt") |
|
|
anno_bboxes = parse_anno(anno_file_path) |
|
|
if anno_bboxes is None: |
|
|
continue |
|
|
|
|
|
cur_video_folder = os.path.join(src_frames_folder, instance_name, 'imgs') |
|
|
frame_names = os.listdir(cur_video_folder) |
|
|
len_frames = len(frame_names) |
|
|
|
|
|
if len_frames > len(anno_bboxes): |
|
|
print(f"Wrong anno and seq, {len_frames} frames, {len(anno_bboxes)} bboxes.") |
|
|
continue |
|
|
|
|
|
print(instance_name) |
|
|
|
|
|
frame_steps = len_frames // _sub_nums |
|
|
for _sub_idx in range(_sub_nums): |
|
|
frame_start_idx = _sub_idx * frame_steps |
|
|
frame_end_idx = min((_sub_idx + 1) * frame_steps, len_frames) |
|
|
|
|
|
selected_frames_idxs = list(range(frame_start_idx, frame_end_idx)) |
|
|
random.shuffle(selected_frames_idxs) |
|
|
selected_frames_idxs = selected_frames_idxs[:_SAMPLE_FRAMES] |
|
|
selected_frames_idxs.sort() |
|
|
|
|
|
if anno_bboxes[selected_frames_idxs[0]] is None: |
|
|
selected_frames_idxs.append(find_nearest(anno_bboxes, selected_frames_idxs[0])) |
|
|
selected_frames_idxs.sort() |
|
|
|
|
|
|
|
|
str_id = str(_id)[1:] |
|
|
_id += 1 |
|
|
drt_folder = os.path.join('./achieved/images/', str_id) |
|
|
if not os.path.exists(drt_folder): |
|
|
os.mkdir(drt_folder) |
|
|
for select_frame_idx in selected_frames_idxs: |
|
|
frame_name = frame_names[select_frame_idx] |
|
|
os.system(f"cp {os.path.join(cur_video_folder, frame_name)} {drt_folder}") |
|
|
|
|
|
|
|
|
selected_anns = [] |
|
|
print(len(anno_bboxes), '--', selected_frames_idxs) |
|
|
for select_frame_idx in selected_frames_idxs: |
|
|
selected_anns.append(anno_bboxes[select_frame_idx]) |
|
|
|
|
|
_data = {"id": "vt_vot{}".format(str_id)} |
|
|
_data["input"] = {"video_folder": drt_folder.replace('/achieved', ''), "prompt": "Please tracking the object within red box in image 1."} |
|
|
_data["output"] = {"bboxes": selected_anns} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
first_frame = Image.open(os.path.join(drt_folder, frame_names[selected_frames_idxs[0]])) |
|
|
draw = ImageDraw.Draw(first_frame) |
|
|
draw.rectangle([selected_anns[0][0], selected_anns[0][1], |
|
|
selected_anns[0][2] + selected_anns[0][0], |
|
|
selected_anns[0][3] + selected_anns[0][1]], outline='red', width=2) |
|
|
first_frame.save(os.path.join(drt_folder, frame_names[selected_frames_idxs[0]].replace('.jpg', '_draw.jpg'))) |
|
|
split_data_list[_split].append(_data) |
|
|
|
|
|
for _split in split_data_list.keys(): |
|
|
with open(f'./achieved/{_split}.json', 'w') as f: |
|
|
print(len(split_data_list[_split])) |
|
|
_data = split_data_list[_split] |
|
|
_ret_data = copy.deepcopy(final_json_data) |
|
|
_ret_data["task"] += f", {_split}" |
|
|
_ret_data["data"] = _data |
|
|
json.dump(_ret_data, f) |
