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/results_obb2hbb.py
Views: 475
1
import os

2
import argparse

3
import shutil

4


5
def parse_args():

6
    parser = argparse.ArgumentParser(description='Train a detector')

7
    parser.add_argument('--srcpath', default=r'/OrientedRepPoints/tools/parse_pkl/evaluation_results/OBB_results/')

8
    parser.add_argument('--dstpath', default='/OrientedRepPoints/tools/parse_pkl/evaluation_results/HBB_results/',

9
                        help='dota version')

10
    args = parser.parse_args()

11


12
    return args

13


14


15
def GetFileFromThisRootDir(dir,ext = None):

16
  allfiles = []

17
  needExtFilter = (ext != None)

18
  for root,dirs,files in os.walk(dir):

19
    for filespath in files:

20
      filepath = os.path.join(root, filespath)

21
      extension = os.path.splitext(filepath)[1][1:]

22
      if needExtFilter and extension in ext:

23
        allfiles.append(filepath)

24
      elif not needExtFilter:

25
        allfiles.append(filepath)

26
  return allfiles

27


28
def custombasename(fullname):

29
    return os.path.basename(os.path.splitext(fullname)[0])

30


31
def OBB2HBB(srcpath, dstpath):

32
    filenames = GetFileFromThisRootDir(srcpath)

33
    if os.path.exists(dstpath):

34
      shutil.rmtree(dstpath)  # delete output folderX

35
    os.makedirs(dstpath)

36


37
    for file in filenames:   # eg: /.../task1_plane.txt

38


39
        basename = custombasename(file)  # 只留下文件名 eg:'task1_plane'

40
        class_basename = basename.split('_')[-1]

41
        with open(file, 'r') as f_in:

42
            with open(os.path.join(dstpath, 'Task2_' + class_basename + '.txt'), 'w') as f_out:

43
                lines = f_in.readlines() 

44
                splitlines = [x.strip().split() for x in lines]  # list: n*[]

45
                for index, splitline in enumerate(splitlines):

46
                    imgname = splitline[0]

47
                    score = splitline[1]

48
                    poly = splitline[2:]

49
                    poly = list(map(float, poly))

50
                    xmin, xmax, ymin, ymax = min(poly[0::2]), max(poly[0::2]), min(poly[1::2]), max(poly[1::2])

51
                    rec_poly = [xmin, ymin, xmax, ymax]

52
                    outline = imgname + ' ' + score + ' ' + ' '.join(map(str, rec_poly))

53
                    if index != (len(splitlines) - 1):

54
                        outline = outline + '\n'

55
                    f_out.write(outline)

56


57
if __name__ == '__main__':

58
    args = parse_args()

59
    OBB2HBB(args.srcpath, args.dstpath)

60


61