RCRS_ML_Model/Labeler_test.py at master · devswha/RCRS_ML_Model · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
from module.myModule import *

mapName = ""
GRID_COL = 0
GRID_ROW = 0

if not len(sys.argv) is 4:
    print("Usage : python Labeler_test.py [Map name][Grid row][Grid col]")
    exit(1)
else:
    mapName = sys.argv[1]
    GRID_ROW = int(sys.argv[2])
    GRID_COL = int(sys.argv[3])

# Label matrix
# [ Number of civilians ]
LABEL = [[0]*GRID_COL for i in range(GRID_ROW)]

# Grid matrix
# [ minX, maxX, minY, maxY ]
GRID = [[[0, 0, 0, 0]]*GRID_COL for i in range(GRID_ROW)]
grid = "%dx%d" % (GRID_ROW, GRID_COL)
if not os.path.exists("%s/%s/%s" % (MODELS_DIR, mapName, grid)):
    os.makedirs("%s/%s/%s" % (MODELS_DIR, mapName, grid))

# Log
now = datetime.datetime.now().strftime("%d %H:%M:%S")
print('[INFO][%s] Image label start' % now)
print('[INFO] Data sets root: %s/%s' % (LABEL_TEST_DATASET_DIR, mapName))
print('[INFO] Label %s #%d ~ #%d data sets' % (mapName, LABEL_TEST_START_MAP_NUM, LABEL_TEST_END_MAP_NUM))

# Filtering and Labeling the images
# Filtering : crop and resize
# Labeling : depend on the time step limit and civilians HP
HPListIndex = 0
for dataSetNum in range(LABEL_TEST_START_MAP_NUM, LABEL_TEST_END_MAP_NUM+1):
    dataSetPath = "%s/raw/test/generated_image/%s/%s_%d" % (LABEL_TEST_DATASET_DIR, mapName, mapName, dataSetNum)

    # Read information files for labeling
    mapInfoFile = open("%s/Parse/mapInfo.txt" % dataSetPath, 'r')
    civilianLocFile = open("%s/Parse/civilianLoc.txt" % dataSetPath, 'r')
    civilianHPFile = open("%s/Parse/civilianHP.txt" % dataSetPath, 'r')

    # Create or Copy the Label directory
    if os.path.exists("%s/Label/%s" % (dataSetPath, grid)):
        shutil.rmtree("%s/Label/%s" % (dataSetPath, grid))
    os.makedirs("%s/Label/%s" % (dataSetPath, grid))

    # Set map's width and height
    initWidth = 0
    initHeight = 0
    endWidth = mapInfoFile.readline().strip('\n')
    endHeight = mapInfoFile.readline().strip('\n')

    # Assign grid cell's range
    for row in range(0, GRID_ROW):
        for col in range(0, GRID_COL):
            GRID[row][col] = [(col * (int(float(endWidth)) / GRID_COL)),
                              ((col + 1) * (int(float(endWidth)) / GRID_COL)),
                              ((GRID_ROW - (row + 1)) * (int(float(endHeight)) / GRID_ROW)),
                              ((GRID_ROW - row) * (int(float(endHeight)) / GRID_ROW))]

    # Labeling the images
    imageListFile = open("%s/Label/%s/ImageList.txt" % (dataSetPath, grid), "w+")
    while True:
        # Parsing the map data
        # Read time step
        line = civilianLocFile.readline().strip('\n')
        if not line: break
        step = int(line)
        civilianHPFile.readline()
        # Read civilian locations and HP list
        HPList = civilianHPFile.readline().split(' ')
        locList = civilianLocFile.readline().split('>')
        for eachLoc in locList:
            eachLoc = eachLoc.strip(' ')
            eachLoc = eachLoc.strip('<')
            if eachLoc != '' and eachLoc != '\n':
                eachLoc = eachLoc.split(', ')
                x = int(eachLoc[0])
                y = int(eachLoc[1])

                # Check civilian's coordinate within cell
                for row in range(0, GRID_ROW):
                    for col in range(0, GRID_COL):
                        if GRID[row][col][0] <= x <= GRID[row][col][1] and GRID[row][col][2] <= y <= GRID[row][col][3]:
                            if int(HPList[HPListIndex]) <= LIMIT_CIVILIAN_HP:
                                LABEL[row][col] = int(LABEL[row][col]) + 1
            HPListIndex = HPListIndex + 1

        # Copy converted image to data set's label directory
        # Longer than LIMIT_TIME_STEP
        if step > LIMIT_TIME_STEP:
            label = []
            for row in range(0, GRID_ROW):
                for col in range(0, GRID_COL):
                    label.append(int(LABEL[row][col]))
            image = "%s/Image/%s_%d_Time_%d.png" % (dataSetPath, mapName, dataSetNum, step)
            if os.path.isfile("%s" % image):
                imageListFile.write(image)
                for eachLabel in label:
                    imageListFile.write(" " + str(eachLabel))
                imageListFile.write("\n")
            else:
                print("[INFO] %s %s is not exist." % (dataSetPath, image))

        # Initialize label matrix
        LABEL = [[0] * GRID_COL for i in range(GRID_ROW)]
        HPListIndex = 0
        now = datetime.datetime.now().strftime("%d %H:%M:%S")
    imageListFile.close()
    print('[INFO][%s] %s_%d images are labeled' % (now, mapName, dataSetNum))