Spaces:
Runtime error
Runtime error
File size: 3,117 Bytes
f670afc |
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 |
# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import os
import csv
class ReducedLabelMapper:
def __init__(self):
this_path = os.path.dirname(os.path.abspath(__file__))
print('[ReducedLabelMapper] Loading from {}'.format(this_path))
# Load Minecraft LUT
mcid2rdlbl_lut = {}
mcid2mclbl_lut = {}
with open(os.path.join(this_path, 'mc_reduction.csv'), newline='') as csvfile:
csvreader = csv.reader(csvfile, delimiter=',')
for row in csvreader:
mcid = int(row[0])
mcid2rdlbl_lut[mcid] = row[3]
mcid2mclbl_lut[mcid] = row[1]
# Load reduced label set
reduced_lbls = []
rdlbl2rdid = {}
with open(os.path.join(this_path, 'reduced_coco_lbls.csv'), newline='') as csvfile:
csvreader = csv.reader(csvfile, delimiter=',')
for idx, row in enumerate(csvreader):
rdlbl2rdid[row[0]] = idx
reduced_lbls.append(row[0])
print(['{}: {}'.format(rdid, rdlbl) for rdid, rdlbl in enumerate(reduced_lbls)])
# The first label should always be 'ignore'
assert reduced_lbls[0] == 'ignore'
# Generate Minecraft ID to Reduced ID LUT
mcid2rdid_lut = []
for mcid in range(len(mcid2rdlbl_lut)):
rdlbl = mcid2rdlbl_lut[mcid]
if rdlbl == '':
rdlbl = 'ignore'
rdid = rdlbl2rdid[rdlbl]
mcid2rdid_lut.append(rdid)
# ================= coco part ==================
gg_label_list = []
gglbl2ggid = {}
with open(os.path.join(this_path, 'gaugan_lbl2col.csv'), newline='') as csvfile:
csvreader = csv.reader(csvfile, delimiter=',')
for idx, row in enumerate(csvreader):
gg_label_list.append(row[0])
gglbl2ggid[row[0]] = idx
# Load coco -> reduced mapping table
gglbl2rdid = {}
with open(os.path.join(this_path, 'gaugan_reduction.csv'), newline='') as csvfile:
csvreader = csv.reader(csvfile, delimiter=',')
for idx, row in enumerate(csvreader):
gglbl = row[0]
target_rdlbl = row[1]
ggid = gglbl2ggid[gglbl]
target_rdid = rdlbl2rdid[target_rdlbl]
gglbl2rdid[ggid] = target_rdid
ggid2rdid = [gglbl2rdid[i] for i in range(len(gglbl2rdid))]
print('[ReducedLabelMapper] #Reduced Labels: {}'.format(len(reduced_lbls)))
self.mcid2rdid_lut = mcid2rdid_lut
self.ggid2rdid = ggid2rdid
self.reduced_lbls = reduced_lbls
self.ignore_id = rdlbl2rdid['ignore']
self.dirt_id = rdlbl2rdid['dirt']
self.water_id = rdlbl2rdid['water']
self.gglbl2ggid = gglbl2ggid
def gglbl2ggid(self, gglbl):
return self.gglbl2ggid[gglbl]
if __name__ == '__main__':
mapper = ReducedLabelMapper()
|