def __init__(self, data_dir, n_triplets, transform=None, transform_parameter=np.array([1.0, 1.0, 1.0]), eval=False, ratio=0.8): super(TripletDataset, self).__init__() self.dir = data_dir self.n_triplets = n_triplets self.transform = transform self.transform_parameter = transform_parameter self.eval = eval with open('LabelEncoder.pkl', 'rb') as f: self.le = pkl.load(f) self.train_csv = pd.read_csv('/data1/whale/train.csv').as_matrix() np.random.shuffle(self.train_csv) self.data = os.listdir(self.dir) print('Generating {} triplets'.format(self.n_triplets)) self.triplets = generate_triplets(self.train_csv, self.le, self.n_triplets, len(self.le.classes_))
#!/usr/bin/env python3 from utils import load_images, load_csv, generate_triplets import numpy as np import matplotlib.pyplot as plt images, filenames = load_images('HBTNaligned', as_array=True) triplet_names = load_csv('FVTriplets.csv') print("triplet_names", len(triplet_names)) A, P, N = generate_triplets(images, filenames, triplet_names) print("A ", A.shape) print("P ", P.shape) print("N ", N.shape) plt.subplot(1, 3, 1) plt.title('Anchor:' + triplet_names[0][0]) plt.imshow(A[0]) plt.subplot(1, 3, 2) plt.title('Positive:' + triplet_names[0][1]) plt.imshow(P[0]) plt.subplot(1, 3, 3) plt.title('Negative:' + triplet_names[0][2]) plt.imshow(N[0]) plt.tight_layout() plt.show()