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

2
import numpy as np

3
import cv2

4
import copy

5
import dota_utils as util

6
from multiprocessing import Pool

7
from functools import partial

8


9


10
def split_single_warp(name, split_base, rate, extent):

11
    split_base.SplitSingle(name, rate, extent)

12


13


14
class splitbase():

15
    def __init__(self,

16
                 srcpath,

17
                 dstpath,

18
                 gap=100,

19
                 subsize=1024,

20
                 ext='.png',

21
                 padding=True,

22
                 num_process=32):

23
        self.srcpath = srcpath

24
        self.outpath = dstpath

25
        self.gap = gap

26
        self.subsize = subsize

27
        self.slide = self.subsize - self.gap

28
        self.srcpath = srcpath

29
        self.dstpath = dstpath

30
        self.ext = ext

31
        self.padding = padding

32
        self.pool = Pool(num_process)

33


34
        if not os.path.isdir(self.outpath):

35
            os.mkdir(self.outpath)

36


37
    def saveimagepatches(self, img, subimgname, left, up, ext='.png'):

38
        subimg = copy.deepcopy(

39
            img[up: (up + self.subsize), left: (left + self.subsize)])

40
        outdir = os.path.join(self.dstpath, subimgname + ext)

41
        h, w, c = np.shape(subimg)

42
        if (self.padding):

43
            outimg = np.zeros((self.subsize, self.subsize, 3))

44
            outimg[0:h, 0:w, :] = subimg

45
            cv2.imwrite(outdir, outimg)

46
        else:

47
            cv2.imwrite(outdir, subimg)

48


49
    def SplitSingle(self, name, rate, extent):

50
        img = cv2.imread(os.path.join(self.srcpath, name + extent))

51
        assert np.shape(img) != ()

52


53
        if (rate != 1):

54
            resizeimg = cv2.resize(

55
                img, None, fx=rate, fy=rate, interpolation=cv2.INTER_CUBIC)

56
        else:

57
            resizeimg = img

58
        outbasename = name + '__' + str(rate) + '__'

59


60
        weight = np.shape(resizeimg)[1]

61
        height = np.shape(resizeimg)[0]

62


63
        # if (max(weight, height) < self.subsize/2):

64
        #     return

65


66
        left, up = 0, 0

67
        while (left < weight):

68
            if (left + self.subsize >= weight):

69
                left = max(weight - self.subsize, 0)

70
            up = 0

71
            while (up < height):

72
                if (up + self.subsize >= height):

73
                    up = max(height - self.subsize, 0)

74
                subimgname = outbasename + str(left) + '___' + str(up)

75
                self.saveimagepatches(resizeimg, subimgname, left, up)

76
                if (up + self.subsize >= height):

77
                    break

78
                else:

79
                    up = up + self.slide

80
            if (left + self.subsize >= weight):

81
                break

82
            else:

83
                left = left + self.slide

84


85
    def splitdata(self, rate):

86


87
        imagelist = util.GetFileFromThisRootDir(self.srcpath)

88
        imagenames = [util.custombasename(x) for x in imagelist if (

89
            util.custombasename(x) != 'Thumbs')]

90


91
        # worker = partial(self.SplitSingle, rate=rate, extent=self.ext)

92
        worker = partial(split_single_warp, split_base=self,

93
                         rate=rate, extent=self.ext)

94
        self.pool.map(worker, imagenames)

95
        #

96
        # for name in imagenames:

97
        #     self.SplitSingle(name, rate, self.ext)

98


99
    def __getstate__(self):

100
        self_dict = self.__dict__.copy()

101
        del self_dict['pool']

102
        return self_dict

103


104
    def __setstate__(self, state):

105
        self.__dict__.update(state)

106


107


108
if __name__ == '__main__':

109
    split = splitbase(r'/media/test/4d846cae-2315-4928-8d1b-ca6d3a61a3c6/DOTA/DOTAv1.5/test/images', 

110
                      r'/media/test/4d846cae-2315-4928-8d1b-ca6d3a61a3c6/DOTA/DOTAv2.0/test-dev_split/images',

111
                      gap=200, subsize=1024, num_process=8)

112
    split.splitdata(1)

113
    # split.splitdata(0.5)

114
    # split.splitdata(1.5)

115
    print("Split Done!")

116


117


118