Update pretreatment_SUSTech1K.py
Fix RGB_image transposed shape
This commit is contained in:
@@ -54,8 +54,8 @@ def imgs2pickle(img_groups: Tuple, output_path: Path, img_size: int = 64, verbos
|
||||
HWs = [img.shape[:2] for img in imgs]
|
||||
# transpose to (C, H W)
|
||||
data = [cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_CUBIC) for img in imgs]
|
||||
imgs = [img.transpose(2, 0, 1) for img in imgs]
|
||||
data = np.asarray(data)
|
||||
imgs = [img.transpose(2, 0, 1) for img in data]
|
||||
data = np.asarray(imgs)
|
||||
HWs = np.asarray(HWs)
|
||||
elif modality == 'Sils_raw':
|
||||
sils = [cv2.imread(sil, cv2.IMREAD_GRAYSCALE) for sil in data_files]
|
||||
@@ -218,4 +218,4 @@ if __name__ == '__main__':
|
||||
for k, v in args.__dict__.items():
|
||||
logging.debug(f'{k}: {v}')
|
||||
|
||||
pretreat(input_path=Path(args.input_path), output_path=Path(args.output_path), img_size=args.img_size, workers=args.n_workers, verbose=args.verbose, dataset=args.dataset)
|
||||
pretreat(input_path=Path(args.input_path), output_path=Path(args.output_path), img_size=args.img_size, workers=args.n_workers, verbose=args.verbose, dataset=args.dataset)
|
||||
|
||||
Reference in New Issue
Block a user