164 lines
5.1 KiB
Python
164 lines
5.1 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# Reference: https://github.com/CAPTAIN-WHU/DOTA_devkit
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import os
|
|
import json
|
|
import cv2
|
|
from tqdm import tqdm
|
|
from multiprocessing import Pool
|
|
|
|
|
|
def load_dota_info(image_dir, anno_dir, file_name, ext=None):
|
|
base_name, extension = os.path.splitext(file_name)
|
|
if ext and (extension != ext and extension not in ext):
|
|
return None
|
|
info = {'image_file': os.path.join(image_dir, file_name), 'annotation': []}
|
|
anno_file = os.path.join(anno_dir, base_name + '.txt')
|
|
if not os.path.exists(anno_file):
|
|
return info
|
|
with open(anno_file, 'r') as f:
|
|
for line in f:
|
|
items = line.strip().split()
|
|
if (len(items) < 9):
|
|
continue
|
|
|
|
anno = {
|
|
'poly': list(map(float, items[:8])),
|
|
'name': items[8],
|
|
'difficult': '0' if len(items) == 9 else items[9],
|
|
}
|
|
info['annotation'].append(anno)
|
|
|
|
return info
|
|
|
|
|
|
def load_dota_infos(root_dir, num_process=8, ext=None):
|
|
image_dir = os.path.join(root_dir, 'images')
|
|
anno_dir = os.path.join(root_dir, 'labelTxt')
|
|
data_infos = []
|
|
if num_process > 1:
|
|
pool = Pool(num_process)
|
|
results = []
|
|
for file_name in os.listdir(image_dir):
|
|
results.append(
|
|
pool.apply_async(load_dota_info, (image_dir, anno_dir,
|
|
file_name, ext)))
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
for result in results:
|
|
info = result.get()
|
|
if info:
|
|
data_infos.append(info)
|
|
|
|
else:
|
|
for file_name in os.listdir(image_dir):
|
|
info = load_dota_info(image_dir, anno_dir, file_name, ext)
|
|
if info:
|
|
data_infos.append(info)
|
|
|
|
return data_infos
|
|
|
|
|
|
def process_single_sample(info, image_id, class_names):
|
|
image_file = info['image_file']
|
|
single_image = dict()
|
|
single_image['file_name'] = os.path.split(image_file)[-1]
|
|
single_image['id'] = image_id
|
|
image = cv2.imread(image_file)
|
|
height, width, _ = image.shape
|
|
single_image['width'] = width
|
|
single_image['height'] = height
|
|
|
|
# process annotation field
|
|
single_objs = []
|
|
objects = info['annotation']
|
|
for obj in objects:
|
|
poly, name, difficult = obj['poly'], obj['name'], obj['difficult']
|
|
if difficult == '2':
|
|
continue
|
|
|
|
single_obj = dict()
|
|
single_obj['category_id'] = class_names.index(name) + 1
|
|
single_obj['segmentation'] = [poly]
|
|
single_obj['iscrowd'] = 0
|
|
xmin, ymin, xmax, ymax = min(poly[0::2]), min(poly[1::2]), max(poly[
|
|
0::2]), max(poly[1::2])
|
|
width, height = xmax - xmin, ymax - ymin
|
|
single_obj['bbox'] = [xmin, ymin, width, height]
|
|
single_obj['area'] = height * width
|
|
single_obj['image_id'] = image_id
|
|
single_objs.append(single_obj)
|
|
|
|
return (single_image, single_objs)
|
|
|
|
|
|
def data_to_coco(infos, output_path, class_names, num_process):
|
|
data_dict = dict()
|
|
data_dict['categories'] = []
|
|
|
|
for i, name in enumerate(class_names):
|
|
data_dict['categories'].append({
|
|
'id': i + 1,
|
|
'name': name,
|
|
'supercategory': name
|
|
})
|
|
|
|
pbar = tqdm(total=len(infos), desc='data to coco')
|
|
images, annotations = [], []
|
|
if num_process > 1:
|
|
pool = Pool(num_process)
|
|
results = []
|
|
for i, info in enumerate(infos):
|
|
image_id = i + 1
|
|
results.append(
|
|
pool.apply_async(
|
|
process_single_sample, (info, image_id, class_names),
|
|
callback=lambda x: pbar.update()))
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
for result in results:
|
|
single_image, single_anno = result.get()
|
|
images.append(single_image)
|
|
annotations += single_anno
|
|
|
|
else:
|
|
for i, info in enumerate(infos):
|
|
image_id = i + 1
|
|
single_image, single_anno = process_single_sample(info, image_id,
|
|
class_names)
|
|
images.append(single_image)
|
|
annotations += single_anno
|
|
pbar.update()
|
|
|
|
pbar.close()
|
|
|
|
for i, anno in enumerate(annotations):
|
|
anno['id'] = i + 1
|
|
|
|
data_dict['images'] = images
|
|
data_dict['annotations'] = annotations
|
|
|
|
with open(output_path, 'w') as f:
|
|
json.dump(data_dict, f)
|