def supply_collate(self, batch):
"""
args:
batch - list of (tensor, label)
reutrn:
xs - a tensor of all examples in 'batch' after padding
ys - a LongTensor of all labels in batch
"""
data = []
target = []
bs = len(batch)
for i in range(bs):
try:
data.append([process_data(np.array(batch[i][0]))])
target.append(batch[i][1])
except:
while True:
try:
new_index = random.randint(0, len(self.dataset) - 1)
data.append([
process_data(np.array(self.dataset[new_index][0]))
])
target.append(self.dataset[new_index][1])
break
except:
continue
data = torch.FloatTensor(data).squeeze(1).permute(0, 3, 1, 2) / 255
target = torch.LongTensor(target)
# print(data.shape, target.shape)
assert len(data) == bs and len(target) == bs
return [data, target]
def my_collate_fn(batch):
"""
detect_align data to form a batch
"""
# print(np.array(batch[0][0]).shape)
data = [process_data(np.array(item[0]))
for item in batch] # No trans version
# data = [process_data((item[0]*255).permute(1,2,0).numpy().astype(np.uint8())) for item in batch]
target = [item[1] for item in batch]
data = torch.FloatTensor(data / 255).permute(0, 2, 3, 1)
target = torch.LongTensor(target)
return [data, target]
from detect_align import process_data
from tqdm import tqdm
import threading
ori_datapath = '/media/a/HDD/Face_Proj/vgg_face_dataset/images'
new_datapath = '/media/a/HDD/Face_Proj/vgg_face_dataset/new_images'
files = os.listdir(ori_datapath)
def download_and_save(path, warped):
cv2.imwrite(path, warped)
for i in tqdm(range(len(files))):
file_name = files[i]
if not os.path.exists(os.path.join(new_datapath, file_name)):
os.makedirs(os.path.join(new_datapath, file_name))
imgs = os.listdir(os.path.join(ori_datapath, file_name))
for j in range(len(imgs)):
img_name = os.path.join(ori_datapath, file_name, imgs[j])
try:
warped = process_data(cv2.imread(img_name))
path = os.path.join(new_datapath, file_name, imgs[j])
t = threading.Thread(target=download_and_save, args=(path, warped))
t.start()
# cv2.imwrite(, warped)
except Exception as e:
continue