kittendev commited on
Commit
845a60a
1 Parent(s): 4619785

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +132 -109
app.py CHANGED
@@ -65,116 +65,139 @@ parser.add_argument('--domain', default='val', type=str)
65
  parser.add_argument('--scales', default='0.5,1.0,1.5,2.0', type=str)
66
  parser.add_argument('--iteration', default=10, type=int)
67
 
68
- dataset = {
69
- "colors": {
70
- "background": [
71
- 0,
72
- 0,
73
- 0
74
- ],
75
- "aeroplane": [
76
- 128,
77
- 0,
78
- 0
79
- ],
80
- "bicycle": [
81
- 0,
82
- 128,
83
- 0
84
- ],
85
- "bird": [
86
- 128,
87
- 128,
88
- 0
89
- ],
90
- "boat": [
91
- 0,
92
- 0,
93
- 128
94
- ],
95
- "bottle": [
96
- 128,
97
- 0,
98
- 128
99
- ],
100
- "bus": [
101
- 0,
102
- 128,
103
- 128
104
- ],
105
- "car": [
106
- 128,
107
- 128,
108
- 128
109
- ],
110
- "cat": [
111
- 64,
112
- 0,
113
- 0
114
- ],
115
- "chair": [
116
- 192,
117
- 0,
118
- 0
119
- ],
120
- "cow": [
121
- 64,
122
- 128,
123
- 0
124
- ],
125
- "diningtable": [
126
- 192,
127
- 128,
128
- 0
129
- ],
130
- "dog": [
131
- 64,
132
- 0,
133
- 128
134
- ],
135
- "horse": [
136
- 192,
137
- 0,
138
- 128
139
- ],
140
- "motorbike": [
141
- 64,
142
- 128,
143
- 128
144
- ],
145
- "person": [
146
- 192,
147
- 128,
148
- 128
149
- ],
150
- "pottedplant": [
151
- 0,
152
- 64,
153
- 0
154
- ],
155
- "sheep": [
156
- 128,
157
- 64,
158
- 0
159
- ],
160
- "sofa": [
161
- 0,
162
- 192,
163
- 0
164
- ],
165
- "train": [
166
- 128,
167
- 192,
168
- 0
169
- ],
170
- "tvmonitor": [
171
- 0,
172
- 64,
173
- 128
174
- ]
175
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
176
  }
177
 
 
 
178
  if __name__ == '__main__':
179
  ###################################################################################
180
  # Arguments
@@ -275,7 +298,7 @@ if __name__ == '__main__':
275
  else:
276
  pred_mask = np.argmax(preds, axis=-1)
277
 
278
- pred_mask = decode_from_colormap(pred_mask, dataset.colors)[..., ::-1]
279
 
280
  return pred_mask.astype(np.uint8)
281
 
 
65
  parser.add_argument('--scales', default='0.5,1.0,1.5,2.0', type=str)
66
  parser.add_argument('--iteration', default=10, type=int)
67
 
68
+ class_names = [
69
+ "aeroplane",
70
+ "bicycle",
71
+ "bird",
72
+ "boat",
73
+ "bottle",
74
+ "bus",
75
+ "car",
76
+ "cat",
77
+ "chair",
78
+ "cow",
79
+ "diningtable",
80
+ "dog",
81
+ "horse",
82
+ "motorbike",
83
+ "person",
84
+ "pottedplant",
85
+ "sheep",
86
+ "sofa",
87
+ "train",
88
+ "tvmonitor"
89
+ ]
90
+
91
+ cmap_dic = {
92
+ "background": [
93
+ 0,
94
+ 0,
95
+ 0
96
+ ],
97
+ "aeroplane": [
98
+ 128,
99
+ 0,
100
+ 0
101
+ ],
102
+ "bicycle": [
103
+ 0,
104
+ 128,
105
+ 0
106
+ ],
107
+ "bird": [
108
+ 128,
109
+ 128,
110
+ 0
111
+ ],
112
+ "boat": [
113
+ 0,
114
+ 0,
115
+ 128
116
+ ],
117
+ "bottle": [
118
+ 128,
119
+ 0,
120
+ 128
121
+ ],
122
+ "bus": [
123
+ 0,
124
+ 128,
125
+ 128
126
+ ],
127
+ "car": [
128
+ 128,
129
+ 128,
130
+ 128
131
+ ],
132
+ "cat": [
133
+ 64,
134
+ 0,
135
+ 0
136
+ ],
137
+ "chair": [
138
+ 192,
139
+ 0,
140
+ 0
141
+ ],
142
+ "cow": [
143
+ 64,
144
+ 128,
145
+ 0
146
+ ],
147
+ "diningtable": [
148
+ 192,
149
+ 128,
150
+ 0
151
+ ],
152
+ "dog": [
153
+ 64,
154
+ 0,
155
+ 128
156
+ ],
157
+ "horse": [
158
+ 192,
159
+ 0,
160
+ 128
161
+ ],
162
+ "motorbike": [
163
+ 64,
164
+ 128,
165
+ 128
166
+ ],
167
+ "person": [
168
+ 192,
169
+ 128,
170
+ 128
171
+ ],
172
+ "pottedplant": [
173
+ 0,
174
+ 64,
175
+ 0
176
+ ],
177
+ "sheep": [
178
+ 128,
179
+ 64,
180
+ 0
181
+ ],
182
+ "sofa": [
183
+ 0,
184
+ 192,
185
+ 0
186
+ ],
187
+ "train": [
188
+ 128,
189
+ 192,
190
+ 0
191
+ ],
192
+ "tvmonitor": [
193
+ 0,
194
+ 64,
195
+ 128
196
+ ]
197
  }
198
 
199
+ colors = np.asarray([cmap_dic[class_name] for class_name in class_names])
200
+
201
  if __name__ == '__main__':
202
  ###################################################################################
203
  # Arguments
 
298
  else:
299
  pred_mask = np.argmax(preds, axis=-1)
300
 
301
+ pred_mask = decode_from_colormap(pred_mask, colors)[..., ::-1]
302
 
303
  return pred_mask.astype(np.uint8)
304