更换文档检测模型

paddle_detection/configs/rotate/tools/convert.py

@@ -0,0 +1,163 @@
+# 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)