plutosss/ImageProcessing
0
1import os
2import random
3
4import cv2
5import numpy as np
6import torch
7from torch.utils.data import Dataset
8import json
9
10DATASET_NAMES = [
11 'BIPED',
12 'BIPED-B2',
13 'BIPED-B3',
14 'BIPED-B5',
15 'BIPED-B6',
16 'BSDS', # 5
17 'BRIND', # 6
18 'ICEDA', #7
19 'BSDS300',
20 'CID', #9
21 'DCD',
22 'MDBD', #11
23 'PASCAL',
24 'NYUD', #13
25 'BIPBRI',
26 'UDED', # 15 just for testing
27 'DMRIR',
28 'CLASSIC'
29] # 8
30# [108, 109.451,112.230,137.86]
31BIPED_mean = [103.939,116.779,123.68,137.86]
32
33def dataset_info(dataset_name, is_linux=True):
34 if is_linux:
35
36 config = {
37 'UDED': {
38 'img_height': 512, # 321
39 'img_width': 512, # 481
40 'train_list': None,
41 'test_list': 'test_pair.lst',
42 'data_dir': '/root/workspace/datasets/UDED', # mean_rgb
43 'yita': 0.5,
44 'mean': [104.007, 116.669, 122.679, 137.86]# [104.007, 116.669, 122.679, 137.86]
45 }, #[98.939,111.779,117.68,137.86]
46 'BSDS': {
47 'img_height': 512, #321
48 'img_width': 512, #481
49 'train_list': 'train_pair.lst',
50 'test_list': 'test_pair.lst',
51 'data_dir': '/root/workspace/datasets/BSDS', # mean_rgb
52 'yita': 0.5,
53 'mean': [104.007, 116.669, 122.679, 137.86]
54 },
55 'BRIND': {
56 'img_height': 512, # 321
57 'img_width': 512, # 481
58 'train_list': 'train_pair_all.lst',
59 # all train_pair_all.lst
60 # less train_pair.lst
61 'test_list': 'test_pair.lst',
62 'data_dir': '/root/workspace/datasets/BRIND', # mean_rgb
63 'yita': 0.5,
64 'mean': [104.007, 116.669, 122.679, 137.86]
65 },
66 'ICEDA': {
67 'img_height': 1024, # 321
68 'img_width': 1408, # 481
69 'train_list': None,
70 'test_list': 'test_pair.lst',
71 'data_dir': '/root/workspace/datasets/ICEDA', # mean_rgb
72 'yita': 0.5,
73 'mean': [104.007, 116.669, 122.679, 137.86]
74 },
75 'BSDS300': {
76 'img_height': 512, #321
77 'img_width': 512, #481
78 'test_list': 'test_pair.lst',
79 'train_list': None,
80 'data_dir': '/root/workspace/datasets/BSDS300', # NIR
81 'yita': 0.5,
82 'mean': [104.007, 116.669, 122.679, 137.86]
83 },
84 'PASCAL': {
85 'img_height': 416, # 375
86 'img_width': 512, #500
87 'test_list': 'test_pair.lst',
88 'train_list': None,
89 'data_dir': '/root/datasets/PASCAL', # mean_rgb
90 'yita': 0.3,
91 'mean': [104.007, 116.669, 122.679, 137.86]
92 },
93 'CID': {
94 'img_height': 512,
95 'img_width': 512,
96 'test_list': 'test_pair.lst',
97 'train_list': None,
98 'data_dir': '/root/datasets/CID', # mean_rgb
99 'yita': 0.3,
100 'mean': [104.007, 116.669, 122.679, 137.86]
101 },
102 'NYUD': {
103 'img_height': 448,#425
104 'img_width': 560,#560
105 'test_list': 'test_pair.lst',
106 'train_list': None,
107 'data_dir': '/root/datasets/NYUD', # mean_rgb
108 'yita': 0.5,
109 'mean': [104.007, 116.669, 122.679, 137.86]
110 },
111 'MDBD': {
112 'img_height': 720,
113 'img_width': 1280,
114 'test_list': 'test_pair.lst',
115 'train_list': 'train_pair.lst',
116 'data_dir': '/root/workspace/datasets/MDBD', # mean_rgb
117 'yita': 0.3,
118 'mean': [104.007, 116.669, 122.679, 137.86]
119 },
120 'BIPED': {
121 'img_height': 720, #720 # 1088
122 'img_width': 1280, # 1280 5 1920
123 'test_list': 'test_pair.lst',
124 'train_list': 'train_pair0.lst', # Base augmentation
125 # 'train_list': 'train_pairB3.lst', # another augmentation
126 # 'train_list': 'train_pairB5.lst', # Last augmentation
127 'data_dir': '/root/workspace/datasets/BIPED', # mean_rgb
128 'yita': 0.5,
129 'mean':BIPED_mean
130 #
131 },
132 'CLASSIC': {
133 'img_height': 512,#
134 'img_width': 512,# 512
135 'test_list': None,
136 'train_list': None,
137 'data_dir': 'data', # mean_rgb
138 'yita': 0.5,
139 'mean': [104.007, 116.669, 122.679, 137.86]
140 },
141 'BIPED-B2': {'img_height': 720, # 720
142 'img_width': 1280, # 1280
143 'test_list': 'test_pair.lst',
144 'train_list': 'train_rgb.lst',
145 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
146 'yita': 0.5,
147 'mean':BIPED_mean},
148 'BIPED-B3': {'img_height': 720, # 720
149 'img_width': 1280, # 1280
150 'test_list': 'test_pair.lst',
151 'train_list': 'train_rgb.lst',
152 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
153 'yita': 0.5,
154 'mean':BIPED_mean},
155 'BIPED-B5': {'img_height': 720, # 720
156 'img_width': 1280, # 1280
157 'test_list': 'test_pair.lst',
158 'train_list': 'train_rgb.lst',
159 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
160 'yita': 0.5,
161 'mean':BIPED_mean},
162 'BIPED-B6': {'img_height': 720, # 720
163 'img_width': 1280, # 1280
164 'test_list': 'test_pair.lst',
165 'train_list': 'train_rgb.lst',
166 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
167 'yita': 0.5,
168 'mean':BIPED_mean},
169 'DCD': {
170 'img_height': 352, #240
171 'img_width': 480,# 360
172 'test_list': 'test_pair.lst',
173 'train_list': None,
174 'data_dir': '/opt/dataset/DCD', # mean_rgb
175 'yita': 0.2,
176 'mean': [104.007, 116.669, 122.679, 137.86]
177 }
178 }
179 else:
180 config = {
181 'UDED': {
182 'img_height': 512, # 321
183 'img_width': 512, # 481
184 'train_list': None,
185 'test_list': 'test_pair.lst',
186 'data_dir': 'C:/dataset/UDED', # mean_rgb
187 'yita': 0.5,
188 'mean':[104.007, 116.669, 122.679, 137.86] # [183.939,196.779,203.68,137.86] # [104.007, 116.669, 122.679, 137.86]
189 },
190 'BSDS': {'img_height': 480, # 321
191 'img_width': 480, # 481
192 'test_list': 'test_pair.lst',
193 'data_dir': 'C:/dataset/BSDS', # mean_rgb
194 'yita': 0.5,
195 'mean':[103.939, 116.669, 122.679, 137.86] },
196 # [103.939, 116.669, 122.679, 137.86]
197 #[159.510, 159.451,162.230,137.86]
198 'BRIND': {
199 'img_height': 512, # 321
200 'img_width': 512, # 481
201 'train_list': 'train_pair_all.lst',
202 # all train_pair_all.lst
203 # less train_pair.lst
204 'test_list': 'test_pair.lst',
205 'data_dir': 'C:/dataset/BRIND', # mean_rgb
206 'yita': 0.5,
207 'mean': [104.007, 116.669, 122.679, 137.86]
208 },
209 'ICEDA': {
210 'img_height': 1024, # 321
211 'img_width': 1408, # 481
212 'train_list': None,
213 'test_list': 'test_pair.lst',
214 'data_dir': 'C:/dataset/ICEDA', # mean_rgb
215 'yita': 0.5,
216 'mean': [104.007, 116.669, 122.679, 137.86]
217 },
218 'BSDS300': {'img_height': 512, # 321
219 'img_width': 512, # 481
220 'test_list': 'test_pair.lst',
221 'data_dir': 'C:/Users/xavysp/dataset/BSDS300', # NIR
222 'yita': 0.5,
223 'mean': [104.007, 116.669, 122.679, 137.86]},
224 'PASCAL': {'img_height': 375,
225 'img_width': 500,
226 'test_list': 'test_pair.lst',
227 'data_dir': 'C:/dataset/PASCAL', # mean_rgb
228 'yita': 0.3,
229 'mean': [104.007, 116.669, 122.679, 137.86]},
230 'CID': {'img_height': 512,
231 'img_width': 512,
232 'test_list': 'test_pair.lst',
233 'data_dir': 'C:/dataset/CID', # mean_rgb
234 'yita': 0.3,
235 'mean': [104.007, 116.669, 122.679, 137.86]},
236 'NYUD': {'img_height': 425,
237 'img_width': 560,
238 'test_list': 'test_pair.lst',
239 'data_dir': 'C:/dataset/NYUD', # mean_rgb
240 'yita': 0.5,
241 'mean': [104.007, 116.669, 122.679, 137.86]},
242 'MDBD': {'img_height': 720,
243 'img_width': 1280,
244 'test_list': 'test_pair.lst',
245 'train_list': 'train_pair.lst',
246 'data_dir': 'C:/dataset/MDBD', # mean_rgb
247 'yita': 0.3,
248 'mean': [104.007, 116.669, 122.679, 137.86]},
249 'BIPED': {'img_height': 720, # 720
250 'img_width': 1280, # 1280
251 'test_list': 'test_pair.lst',
252 'train_list': 'train_pair0.lst',
253 # 'train_list': 'train_rgb.lst',
254 'data_dir': 'C:/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
255 'yita': 0.5,
256 'mean':BIPED_mean},
257 'BIPED-B2': {'img_height': 720, # 720
258 'img_width': 1280, # 1280
259 'test_list': 'test_pair.lst',
260 'train_list': 'train_rgb.lst',
261 'data_dir': 'C:/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
262 'yita': 0.5,
263 'mean':BIPED_mean},
264 'BIPED-B3': {'img_height': 720, # 720
265 'img_width': 1280, # 1280
266 'test_list': 'test_pair.lst',
267 'train_list': 'train_rgb.lst',
268 'data_dir': 'C:/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
269 'yita': 0.5,
270 'mean':BIPED_mean},
271 'BIPED-B5': {'img_height': 720, # 720
272 'img_width': 1280, # 1280
273 'test_list': 'test_pair.lst',
274 'train_list': 'train_rgb.lst',
275 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
276 'yita': 0.5,
277 'mean':BIPED_mean},
278 'BIPED-B6': {'img_height': 720, # 720
279 'img_width': 1280, # 1280
280 'test_list': 'test_pair.lst',
281 'train_list': 'train_rgb.lst',
282 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges'
283 'yita': 0.5,
284 'mean':BIPED_mean},
285 'CLASSIC': {'img_height': 512,
286 'img_width': 512,
287 'test_list': None,
288 'train_list': None,
289 'data_dir': 'teed_tmp', # mean_rgb
290 'yita': 0.5,
291 'mean': [104.007, 116.669, 122.679, 137.86]},
292 'DCD': {'img_height': 240,
293 'img_width': 360,
294 'test_list': 'test_pair.lst',
295 'data_dir': 'C:/dataset/DCD', # mean_rgb
296 'yita': 0.2,
297 'mean': [104.007, 116.669, 122.679, 137.86]}
298 }
299 return config[dataset_name]
300
301class TestDataset(Dataset):
302 def __init__(self,
303 data_root,
304 test_data,
305 img_height,
306 img_width,
307 test_list=None,
308 arg=None
309 ):
310 if test_data not in DATASET_NAMES:
311 raise ValueError(f"Unsupported dataset: {test_data}")
312
313 self.data_root = data_root
314 self.test_data = test_data
315 self.test_list = test_list
316 self.args = arg
317 self.up_scale = arg.up_scale
318 self.mean_bgr = arg.mean_test if len(arg.mean_test) == 3 else arg.mean_test[:3]
319 self.img_height = img_height
320 self.img_width = img_width
321 self.data_index = self._build_index()
322
323
324 def _build_index(self):
325 sample_indices = []
326 if self.test_data == "CLASSIC":
327 # for single image testing
328 images_path = os.listdir(self.data_root)
329 labels_path = None
330 sample_indices = [images_path, labels_path]
331 else:
332 # image and label paths are located in a list file
333
334 if not self.test_list:
335 raise ValueError(
336 f"Test list not provided for dataset: {self.test_data}")
337
338 list_name = os.path.join(self.data_root, self.test_list)
339 if self.test_data.upper() in ['BIPED', 'BRIND','UDED','ICEDA']:
340
341 with open(list_name,encoding='utf-8') as f:
342 files = json.load(f)
343 for pair in files:
344 tmp_img = pair[0]
345 tmp_gt = pair[1]
346 sample_indices.append(
347 (os.path.join(self.data_root, tmp_img),
348 os.path.join(self.data_root, tmp_gt),))
349 else:
350 with open(list_name, 'r') as f:
351 files = f.readlines()
352 files = [line.strip() for line in files]
353 pairs = [line.split() for line in files]
354
355 for pair in pairs:
356 tmp_img = pair[0]
357 tmp_gt = pair[1]
358 sample_indices.append(
359 (os.path.join(self.data_root, tmp_img),
360 os.path.join(self.data_root, tmp_gt),))
361 return sample_indices
362
363 def __len__(self):
364 return len(self.data_index[0]) if self.test_data.upper() == 'CLASSIC' else len(self.data_index)
365
366 def __getitem__(self, idx):
367 # get data sample
368 # image_path, label_path = self.data_index[idx]
369 if self.data_index[1] is None:
370 image_path = self.data_index[0][idx] if len(self.data_index[0]) > 1 else self.data_index[0][idx - 1]
371 else:
372 image_path = self.data_index[idx][0]
373 label_path = None if self.test_data == "CLASSIC" else self.data_index[idx][1]
374 img_name = os.path.basename(image_path)
375 # print(img_name)
376 file_name = os.path.splitext(img_name)[0] + ".png"
377
378 # base dir
379 if self.test_data.upper() == 'BIPED':
380 img_dir = os.path.join(self.data_root, 'imgs', 'test')
381 gt_dir = os.path.join(self.data_root, 'edge_maps', 'test')
382 elif self.test_data.upper() == 'CLASSIC':
383 img_dir = self.data_root
384 gt_dir = None
385 else:
386 img_dir = self.data_root
387 gt_dir = self.data_root
388
389 # load data
390 image = cv2.imdecode(np.fromfile(os.path.join(img_dir, image_path), np.uint8), cv2.IMREAD_COLOR)
391 if not self.test_data == "CLASSIC":
392 label = cv2.imread(os.path.join(
393 gt_dir, label_path), cv2.IMREAD_COLOR)
394 else:
395 label = None
396
397 im_shape = [image.shape[0], image.shape[1]]
398 image, label = self.transform(img=image, gt=label)
399
400 return dict(images=image, labels=label, file_names=file_name, image_shape=im_shape)
401
402 def transform(self, img, gt):
403 # gt[gt< 51] = 0 # test without gt discrimination
404 # up scale test image
405 if self.up_scale:
406 # For TEED BIPBRIlight Upscale
407 img = cv2.resize(img,(0,0),fx=1.3,fy=1.3)
408
409 if img.shape[0] < 512 or img.shape[1] < 512:
410 #TEED BIPED standard proposal if you want speed up the test, comment this block
411 img = cv2.resize(img, (0, 0), fx=1.5, fy=1.5)
412 # else:
413 # img = cv2.resize(img, (0, 0), fx=1.1, fy=1.1)
414
415 # Make sure images and labels are divisible by 2^4=16
416 if img.shape[0] % 8 != 0 or img.shape[1] % 8 != 0:
417 img_width = ((img.shape[1] // 8) + 1) * 8
418 img_height = ((img.shape[0] // 8) + 1) * 8
419 img = cv2.resize(img, (img_width, img_height))
420 # gt = cv2.resize(gt, (img_width, img_height))
421 else:
422 pass
423 # img_width = self.args.test_img_width
424 # img_height = self.args.test_img_height
425 # img = cv2.resize(img, (img_width, img_height))
426 # gt = cv2.resize(gt, (img_width, img_height))
427 # # For FPS
428 # img = cv2.resize(img, (496,320))
429
430 img = np.array(img, dtype=np.float32)
431 # if self.rgb:
432 # img = img[:, :, ::-1] # RGB->BGR
433
434 img -= self.mean_bgr
435 img = img.transpose((2, 0, 1))
436 img = torch.from_numpy(img.copy()).float()
437
438 if self.test_data == "CLASSIC":
439 gt = np.zeros((img.shape[:2]))
440 gt = torch.from_numpy(np.array([gt])).float()
441 else:
442 gt = np.array(gt, dtype=np.float32)
443 if len(gt.shape) == 3:
444 gt = gt[:, :, 0]
445 gt /= 255.
446 gt = torch.from_numpy(np.array([gt])).float()
447
448 return img, gt
449
450# *************************************************
451# ************* training **************************
452# *************************************************
453class BipedDataset(Dataset):
454 train_modes = ['train', 'test', ]
455 dataset_types = ['rgbr', ]
456 data_types = ['aug', ]
457
458 def __init__(self,
459 data_root,
460 img_height,
461 img_width,
462 train_mode='train',
463 dataset_type='rgbr',
464 # is_scaling=None,
465 # Whether to crop image or otherwise resize image to match image height and width.
466 crop_img=False,
467 arg=None
468 ):
469 self.data_root = data_root
470 self.train_mode = train_mode
471 self.dataset_type = dataset_type
472 self.data_type = 'aug' # be aware that this might change in the future
473 self.img_height = img_height
474 self.img_width = img_width
475 self.mean_bgr = arg.mean_train if len(arg.mean_train) == 3 else arg.mean_train[:3]
476 self.crop_img = crop_img
477 self.arg = arg
478
479 self.data_index = self._build_index()
480
481 def _build_index(self):
482 assert self.train_mode in self.train_modes, self.train_mode
483 assert self.dataset_type in self.dataset_types, self.dataset_type
484 assert self.data_type in self.data_types, self.data_type
485
486 data_root = os.path.abspath(self.data_root)
487 sample_indices = []
488
489 file_path = os.path.join(data_root, self.arg.train_list)
490 if self.arg.train_data.lower() == 'bsds':
491
492 with open(file_path, 'r') as f:
493 files = f.readlines()
494 files = [line.strip() for line in files]
495
496 pairs = [line.split() for line in files]
497 for pair in pairs:
498 tmp_img = pair[0]
499 tmp_gt = pair[1]
500 sample_indices.append(
501 (os.path.join(data_root, tmp_img),
502 os.path.join(data_root, tmp_gt),))
503 else:
504 with open(file_path) as f:
505 files = json.load(f)
506 for pair in files:
507 tmp_img = pair[0]
508 tmp_gt = pair[1]
509 sample_indices.append(
510 (os.path.join(data_root, tmp_img),
511 os.path.join(data_root, tmp_gt),))
512
513 return sample_indices
514
515 def __len__(self):
516 return len(self.data_index)
517
518 def __getitem__(self, idx):
519 # get data sample
520 image_path, label_path = self.data_index[idx]
521
522 # load data
523 image = cv2.imdecode(np.fromfile(image_path, np.uint8), cv2.IMREAD_COLOR)
524 label = cv2.imdecode(np.fromfile(label_path), cv2.IMREAD_GRAYSCALE)
525 image, label = self.transform(img=image, gt=label)
526 return dict(images=image, labels=label)
527
528 def transform(self, img, gt):
529 gt = np.array(gt, dtype=np.float32)
530 if len(gt.shape) == 3:
531 gt = gt[:, :, 0]
532
533 gt /= 255. # for LDC input and BDCN
534
535 img = np.array(img, dtype=np.float32)
536 img -= self.mean_bgr
537 i_h, i_w, _ = img.shape
538 # 400 for BIPEd and 352 for BSDS check with 384
539 crop_size = self.img_height if self.img_height == self.img_width else None # 448# MDBD=480 BIPED=480/400 BSDS=352
540 #
541 # # for BSDS 352/BRIND
542 # if i_w > crop_size and i_h > crop_size: # later 400, before crop_size
543 # i = random.randint(0, i_h - crop_size)
544 # j = random.randint(0, i_w - crop_size)
545 # img = img[i:i + crop_size, j:j + crop_size]
546 # gt = gt[i:i + crop_size, j:j + crop_size]
547
548 # for BIPED/MDBD
549 # Second augmentation
550 if i_w> 400 and i_h>400: #before 420
551 h,w = gt.shape
552 if np.random.random() > 0.4: #before i_w> 500 and i_h>500:
553
554 LR_img_size = crop_size #l BIPED=256, 240 200 # MDBD= 352 BSDS= 176
555 i = random.randint(0, h - LR_img_size)
556 j = random.randint(0, w - LR_img_size)
557 # if img.
558 img = img[i:i + LR_img_size , j:j + LR_img_size ]
559 gt = gt[i:i + LR_img_size , j:j + LR_img_size ]
560 else:
561 LR_img_size = 300# 256 300 400 # l BIPED=208-352, # MDBD= 352-480- BSDS= 176-320
562 i = random.randint(0, h - LR_img_size)
563 j = random.randint(0, w - LR_img_size)
564 # if img.
565 img = img[i:i + LR_img_size, j:j + LR_img_size]
566 gt = gt[i:i + LR_img_size, j:j + LR_img_size]
567 img = cv2.resize(img, dsize=(crop_size, crop_size), )
568 gt = cv2.resize(gt, dsize=(crop_size, crop_size))
569
570 else:
571 # New addidings
572 img = cv2.resize(img, dsize=(crop_size, crop_size))
573 gt = cv2.resize(gt, dsize=(crop_size, crop_size))
574 # BRIND Best for TEDD+BIPED
575 gt[gt > 0.1] +=0.2#0.4
576 gt = np.clip(gt, 0., 1.)
577
578 img = img.transpose((2, 0, 1))
579 img = torch.from_numpy(img.copy()).float()
580 gt = torch.from_numpy(np.array([gt])).float()
581 return img, gt