|
- import os
-
- import cv2
- import numpy as np
- import pandas as pd
- import glob
- # from tqdm import tqdm
-
- import data.transforms as transforms
-
- # Per-channel mean and SD values in BGR order
- _MEAN = [0.406, 0.456, 0.485]
- _SD = [0.225, 0.224, 0.229]
-
- # Eig vals and vecs of the cov mat
- _EIG_VALS = np.array([[0.2175, 0.0188, 0.0045]])
- _EIG_VECS = np.array(
- [[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140], [-0.5836, -0.6948, 0.4203]]
- )
-
- class DataLoader:
- def __init__(self, args, info_path):
- self.info_path = info_path
- self.data_root = args.data_url
- self.image_dir = args.data_url + "*/*/*/"
- self.image_paths, self.labels = self.get_clean_train_image_files_and_labels(self.info_path, self.image_dir)
-
- def prepare_image(self, image):
- # Scale and aspect ratio then horizontal flip
- image = transforms.random_sized_crop(im=image, size=321, area_frac=0.08)
- image = transforms.horizontal_flip(im=image, p=0.5, order="HWC")
- # HWC -> CHW
- image = image.transpose([2, 0, 1])
- # [0, 255] -> [0, 1]
- image = image / 255.0
- # PCA jitter
- image = transforms.lighting(image, 0.1, _EIG_VALS, _EIG_VECS)
- # Color normalization
- image = transforms.color_norm(image, _MEAN, _SD)
- return image
-
- def get_data(self, image_path):
- image = cv2.imread(image_path)
- image = image.astype(np.float32, copy=False)
- image_array = self.prepare_image(image)
- return image_array
-
- def get_clean_train_image_files_and_labels(self, csv_path, image_dir):
- """Get image file paths, image ids and labels for the clean training split.
-
- Args:
- csv_path: path to the Google-landmark Dataset v2 CSV Data Sources files
- of the clean train dataset. Assumes CSV header landmark_id;images.
- image_dir: directory that stores downloaded images.
-
- Returns:
- image_paths: the paths to all images in the image_dir.
- file_ids: the unique ids of images.
- labels: the landmark id of all images.
- relabeling: relabeling rules created to replace actual labels with
- a continuous set of labels.
- """
- # Load the content of the CSV file (landmark_id/label -> images).
- csv_file = open(csv_path, 'rb')
- df = pd.read_csv(csv_file)
-
- # Create the dictionary (key = file_id, value = {label, file_id}).
- images = {}
- for _, row in df.iterrows():
- label = row['landmark_id']
- for file_id in row['images'].split(' '):
- images[file_id] = {}
- images[file_id]['label'] = label
- images[file_id]['file_id'] = file_id
-
- # Add the full image path to the dictionary of images.
- image_paths = glob.glob(os.path.join(image_dir, '*.jpg'))
- for image_path in image_paths:
- file_id = os.path.basename(os.path.normpath(image_path))[:-4]
- if file_id in images:
- images[file_id]['image_path'] = image_path.split("data/")[1]
-
- # Explode the dictionary into lists (1 per image attribute).
- image_paths = []
- file_ids = []
- labels = []
- for _, value in images.items():
- try:
- image_paths.append(value['image_path'])
- file_ids.append(value['file_id'])
- labels.append(value['label'])
- except:
- print(value)
- # pass
-
- # Relabel image labels to contiguous values.
- unique_labels = sorted(set(labels))
- relabeling = {label: index for index, label in enumerate(unique_labels)}
- new_labels = [relabeling[label] for label in labels]
-
- # data = pd.DataFrame()
- # data["labels"] = pd.Series(new_labels)
- # data["image_paths"] = pd.Series(image_paths)
-
- return image_paths, new_labels
-
- def get_image_path_and_label(self, info_path):
- image_path_list = []
- label_list = []
- with open(info_path) as fin:
- for line in fin:
- image_dir, label = line.strip().split(" ")
- image_path = self.data_path + image_dir
- if os.path.exists(image_path):
- image_path_list.append(image_path)
- label = np.array([int(label)])
- label_list.append(label)
- else:
- pass
- return image_path_list, label_list
-
- def __getitem__(self, index):
- full_image_path = self.data_root + self.image_paths[index]
- # image_data = self.get_data(full_image_path)
- image_data = cv2.imread(full_image_path)
- return image_data, self.labels[index]
-
- def __len__(self):
- return len(self.image_paths)
|