CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutSign UpSign In
hukaixuan19970627

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

GitHub Repository: hukaixuan19970627/yolov5_obb
 Path: blob/master/DOTA_devkit/DOTA2JSON.py
Views: 475
1
import json

2
import os

3
import os.path as osp

4
import random

5


6
from PIL import Image

7


8
from dota_poly2rbox import poly2rbox_single_v2, poly2rbox_single

9


10


11
def parse_ann_info(img_base_path, label_base_path, img_name):

12
    lab_path = osp.join(label_base_path, img_name+'.txt')

13
    bboxes, labels, bboxes_ignore, labels_ignore = [], [], [], []

14
    with open(lab_path, 'r') as f:

15
        for ann_line in f.readlines():

16
            ann_line = ann_line.strip().split(' ')

17
            bbox = [float(ann_line[i]) for i in range(8)]

18
            # 8 point to 5 point xywha

19
            bbox = poly2rbox_single_v2(bbox)

20
            class_name = ann_line[8]

21
            difficult = int(ann_line[9])

22
            # ignore difficult =2

23
            if difficult == 0:

24
                bboxes.append(bbox)

25
                labels.append(class_name)

26
            elif difficult == 1:

27
                bboxes_ignore.append(bbox)

28
                labels_ignore.append(class_name)

29
    return bboxes, labels, bboxes_ignore, labels_ignore

30


31


32
def generate_txt_labels(src_path, out_path, trainval=True):

33
    """Generate .txt labels recording img_names

34
    Args:

35
        src_path: dataset path containing images and labelTxt folders.

36
        out_path: output txt file path

37
        trainval: trainval or test?

38
    """

39
    img_path = os.path.join(src_path, 'images')

40
    label_path = os.path.join(src_path, 'labelTxt')

41
    img_lists = os.listdir(img_path)

42
    with open(out_path, 'w') as f:

43
        for img in img_lists:

44
            img_name = osp.splitext(img)[0]

45
            label = os.path.join(label_path, img_name+'.txt')

46
            if(trainval == True):

47
                if(os.path.exists(label) == False):

48
                    print('Label:'+img_name+'.txt'+' Not Exist')

49
                else:

50
                    f.write(img_name+'\n')

51
            else:

52
                f.write(img_name+'\n')

53


54


55
def generate_json_labels(src_path, out_path, trainval=True):

56
    """Generate .json labels which is similar to coco format

57
    Args:

58
        src_path: dataset path containing images and labelTxt folders.

59
        out_path: output json file path

60
        trainval: trainval or test?

61
    """

62
    img_path = os.path.join(src_path, 'images')

63
    label_path = os.path.join(src_path, 'labelTxt')

64
    img_lists = os.listdir(img_path)

65


66
    data_dict = []

67


68
    with open(out_path, 'w') as f:

69
        for id, img in enumerate(img_lists):

70
            img_info = {}

71
            img_name = osp.splitext(img)[0]

72
            label = os.path.join(label_path, img_name+'.txt')

73
            img = Image.open(osp.join(img_path, img))

74
            img_info['filename'] = img_name+'.png'

75
            img_info['height'] = img.height

76
            img_info['width'] = img.width

77
            img_info['id'] = id

78
            if(trainval == True):

79
                if(os.path.exists(label) == False):

80
                    print('Label:'+img_name+'.txt'+' Not Exist')

81
                else:

82
                    bboxes, labels, bboxes_ignore, labels_ignore = parse_ann_info(

83
                        img_path, label_path, img_name)

84
                    ann = {}

85
                    ann['bboxes'] = bboxes

86
                    ann['labels'] = labels

87
                    ann['bboxes_ignore'] = bboxes_ignore

88
                    ann['labels_ignore'] = labels_ignore

89
                    img_info['annotations'] = ann

90
            data_dict.append(img_info)

91
        json.dump(data_dict, f)

92


93


94
if __name__ == '__main__':

95
    generate_json_labels('/data1/dataset_demo/DOTA_demo/',

96
                         '/data1/OrientedRepPoints/data/dota/trainval_split/trainval.json')

97
    generate_json_labels('/data1/dataset_demo/DOTA_demo/',

98
                         '/data1/OrientedRepPoints/data/dota/test_split/test.json', trainval=False)

99
    print('done!')

100


101