37 lines
No EOL
954 B
Python
37 lines
No EOL
954 B
Python
import glob
|
|
|
|
import cv2
|
|
|
|
import torchvision.transforms as T
|
|
import torch
|
|
import os
|
|
from tqdm import tqdm
|
|
|
|
|
|
PATH_IN = "/scratch/bortoletto/data/tbd/images"
|
|
PATH_OUT = "/scratch/bortoletto/data/tbd/images_norm"
|
|
|
|
normalisation_steps = [
|
|
T.ToTensor(),
|
|
T.Resize((128,128)),
|
|
T.Normalize(
|
|
mean=[0.485, 0.456, 0.406],
|
|
std=[0.229, 0.224, 0.225]
|
|
)
|
|
]
|
|
|
|
preprocess_img = T.Compose(normalisation_steps)
|
|
|
|
def main():
|
|
print(f"{PATH_IN}/*/*/*.jpg")
|
|
all_img = glob.glob(f"{PATH_IN}/*/*/*.jpg")
|
|
print(len(all_img))
|
|
for img_path in tqdm(all_img):
|
|
new_img = preprocess_img(cv2.imread(img_path)).numpy()
|
|
img_path_split = img_path.split("/")
|
|
os.makedirs(f"{PATH_OUT}/{img_path_split[-3]}/{img_path_split[-2]}", exist_ok=True)
|
|
out_img = f"{PATH_OUT}/{img_path_split[-3]}/{img_path_split[-2]}/{img_path_split[-1][:-4]}.pt"
|
|
torch.save(new_img, out_img)
|
|
|
|
if __name__ == '__main__':
|
|
main() |