136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
|
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
|
|||
|
|
@Time : 11:01 AM
|
|||
|
|
@Author : Zhao kunlong
|
|||
|
|
@Description:
|
|||
|
|
|
|||
|
|
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import cv2 as cv
|
|||
|
|
from glob import glob
|
|||
|
|
from collections import defaultdict
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
CLASSES = ('uav',)
|
|||
|
|
|
|||
|
|
CLASSES_ENCODES = ('None',)
|
|||
|
|
|
|||
|
|
cats_id_maps = {}
|
|||
|
|
for k, v in enumerate(CLASSES_ENCODES, 1):
|
|||
|
|
cats_id_maps[v] = k
|
|||
|
|
|
|||
|
|
|
|||
|
|
def decode_gt(gts: str):
|
|||
|
|
ir_gt = gts.split('\n')[0].split(',')
|
|||
|
|
obj_num = int(ir_gt[1])
|
|||
|
|
frame_id = int(ir_gt[0])
|
|||
|
|
ir_gt_dict = {}
|
|||
|
|
for j in range(obj_num):
|
|||
|
|
obj_id = int(ir_gt[2 + j * 6])
|
|||
|
|
ir_gt_dict[obj_id] = [float(v) for v in ir_gt[3 + j * 6:7 + j * 6]]
|
|||
|
|
|
|||
|
|
return ir_gt_dict, frame_id
|
|||
|
|
|
|||
|
|
def conver2coco(root, VID:defaultdict(list), split:str='train'):
|
|||
|
|
records = dict(
|
|||
|
|
vid_id=1,
|
|||
|
|
img_id=1,
|
|||
|
|
ann_id=1,
|
|||
|
|
global_instance_id=1,
|
|||
|
|
num_vid_train_frames=0,
|
|||
|
|
num_no_objects=0)
|
|||
|
|
|
|||
|
|
trained_frames = 0
|
|||
|
|
videos = glob(os.path.join(root, split,'*'))
|
|||
|
|
pdr = tqdm(videos, desc='processing...')
|
|||
|
|
for _video in pdr:
|
|||
|
|
video_name = os.path.join(*_video.split('/')[-3:])
|
|||
|
|
|
|||
|
|
vid_train_frames = []
|
|||
|
|
|
|||
|
|
gt_path = os.path.join(_video, 'rgb_gt.txt')
|
|||
|
|
with open(gt_path, 'r') as f:
|
|||
|
|
rgb_gts = f.readlines()
|
|||
|
|
rgb_gt_dicts = {}
|
|||
|
|
for rgb_gt in rgb_gts:
|
|||
|
|
rgb_gt_dict, frame_id = decode_gt(rgb_gt)
|
|||
|
|
rgb_gt_dicts[frame_id] = rgb_gt_dict
|
|||
|
|
if len(rgb_gt_dict) != 0:
|
|||
|
|
vid_train_frames.append(frame_id)
|
|||
|
|
|
|||
|
|
trained_frames += len(vid_train_frames)
|
|||
|
|
|
|||
|
|
if split in ['val', 'val_train', 'test']:
|
|||
|
|
video = dict(
|
|||
|
|
id=records['vid_id'],
|
|||
|
|
name=video_name,
|
|||
|
|
vid_train_frames=[])
|
|||
|
|
else:
|
|||
|
|
video = dict(
|
|||
|
|
id=records['vid_id'],
|
|||
|
|
name=video_name,
|
|||
|
|
vid_train_frames=vid_train_frames)
|
|||
|
|
VID['videos'].append(video)
|
|||
|
|
|
|||
|
|
for frame_id in vid_train_frames:
|
|||
|
|
image_path = os.path.join(_video, 'rgb', '{:06d}.jpg'.format(frame_id))
|
|||
|
|
try:
|
|||
|
|
img = cv.imread(image_path)
|
|||
|
|
size = img.shape
|
|||
|
|
height, width, _ = size
|
|||
|
|
except Exception as e:
|
|||
|
|
print(e)
|
|||
|
|
print(image_path)
|
|||
|
|
|
|||
|
|
image = dict(
|
|||
|
|
file_name=os.path.join(*image_path.split('/')[-5:]),
|
|||
|
|
height=height,
|
|||
|
|
width=width,
|
|||
|
|
id=records['img_id'],
|
|||
|
|
frame_id=frame_id,
|
|||
|
|
video_id=records['vid_id'],
|
|||
|
|
is_vid_train_frame=True)
|
|||
|
|
VID['images'].append(image)
|
|||
|
|
|
|||
|
|
if split != 'test ':
|
|||
|
|
obj_annos = rgb_gt_dicts[frame_id]
|
|||
|
|
|
|||
|
|
for obj_anno in obj_annos.values():
|
|||
|
|
x1, y1, x2, y2 = obj_anno
|
|||
|
|
w, h = x2 - x1, y2 - y1
|
|||
|
|
x1, y1, w, h = int(x1), int(y1), int(w), int(h)
|
|||
|
|
|
|||
|
|
ann = dict(
|
|||
|
|
id=records['ann_id'],
|
|||
|
|
video_id=records['vid_id'],
|
|||
|
|
image_id=records['img_id'],
|
|||
|
|
category_id=1,
|
|||
|
|
instance_id=-1,
|
|||
|
|
bbox=[x1, y1, w, h],
|
|||
|
|
area=w * h,
|
|||
|
|
iscrowd=False,
|
|||
|
|
occluded=False,
|
|||
|
|
generated=False)
|
|||
|
|
VID['annotations'].append(ann)
|
|||
|
|
records['ann_id'] += 1
|
|||
|
|
records['img_id'] += 1
|
|||
|
|
records['vid_id'] += 1
|
|||
|
|
assert trained_frames == records['img_id'] - 1
|
|||
|
|
with open(f'sky_data_vid_{split}.json', 'w') as f:
|
|||
|
|
json.dump(VID, f)
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
categories = []
|
|||
|
|
for k, v in enumerate(CLASSES, 1):
|
|||
|
|
categories.append(
|
|||
|
|
dict(id=k, name=v, encode_name=CLASSES_ENCODES[k - 1]))
|
|||
|
|
|
|||
|
|
VID_train = defaultdict(list)
|
|||
|
|
VID_train['categories'] = categories
|
|||
|
|
# root = '/home/baode/project/data/sky_data1'
|
|||
|
|
root = '/home/fty/Documents/fty/mvsod/data/sky_data3'
|
|||
|
|
|
|||
|
|
# conver2coco(root, VID_train, split='train')
|
|||
|
|
conver2coco(root, VID_train, split='val')
|