From 28cdd1245c893df408919b7c6be9572095e8cf94 Mon Sep 17 00:00:00 2001 From: PrimeshES Date: Fri, 4 Feb 2022 16:37:55 +0530 Subject: [PATCH] face3d dataloader updated --- dataloader/face3d.py | 32 ++++++++++++++++++++++++++++++++ test_face3d.py | 12 ++++++------ 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/dataloader/face3d.py b/dataloader/face3d.py index c83ac64..c29f267 100644 --- a/dataloader/face3d.py +++ b/dataloader/face3d.py @@ -519,6 +519,18 @@ def __getitem__(self, idx): gt_label = np.array([gaze_x, gaze_y]) head_box = np.array([x_min/640, y_min/480, x_max/640, y_max/480]) + #plot gaze point + # i = cv2.imread(image_path) + # i=cv2.resize(i, (448, 448)) + # p=(gt_label*448).astype(np.int) + # p2=(head*448).astype(np.int) + # x,y=p + # x2, y2=p2 + # i = cv2.circle(i, (x,y), radius=0, color=(0, 0, 255), thickness=4) + # i = cv2.circle(i, (x2,y2), radius=0, color=(0, 0, 255), thickness=8) + # cv2.imwrite('/home/primesh/Desktop/fig.jpg', i) + # face.save('/home/primesh/Desktop/face.jpg') + if self.imshow: img.save("img_aug.jpg") @@ -593,6 +605,14 @@ def __getitem__(self, idx): x_max = 0 if y_max < 0: y_max = 0 + if x_min > 1: + x_min = 1 + if y_min > 1: + y_min = 1 + if x_max > 1: + x_max = 1 + if y_max > 1: + y_max = 1 x_min -= k * abs(x_max - x_min) y_min -= k * abs(y_max - y_min) x_max += k * abs(x_max - x_min) @@ -610,6 +630,18 @@ def __getitem__(self, idx): gt_label = np.array([gaze_x, gaze_y]) head_box = np.array([x_min/640, y_min/480, x_max/640, y_max/480]) + #plot gaze point + # i = cv2.imread(image_path) + # i=cv2.resize(i, (448, 448)) + # p=(gt_label*448).astype(np.int) + # p2=(head*448).astype(np.int) + # x,y=p + # x2, y2=p2 + # i = cv2.circle(i, (x,y), radius=0, color=(0, 0, 255), thickness=4) + # i = cv2.circle(i, (x2,y2), radius=0, color=(0, 0, 255), thickness=8) + # cv2.imwrite('/home/primesh/Desktop/fig.jpg', i) + # face.save('/home/primesh/Desktop/face.jpg') + if self.imshow: img.save("img_aug.jpg") diff --git a/test_face3d.py b/test_face3d.py index b6ef908..b7fd32b 100644 --- a/test_face3d.py +++ b/test_face3d.py @@ -18,7 +18,7 @@ # from models.__init__ import save_checkpoint, resume_checkpoint from models.face3d import Face3D -from dataloader.face3d import RetailGaze +from dataloader.face3d import RetailGaze, RetailGaze2 from dataloader import chong_imutils from training.face3d import train_face3d, GazeOptimizer, test_face3d from torch.nn.utils.rnn import pad_sequence @@ -41,19 +41,19 @@ print ('Train') -train_set = RetailGaze(images_dir, pickle_path, 'train') +train_set = RetailGaze2(images_dir, pickle_path, 'train') train_data_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True, num_workers=16) print ('Val') -val_set = RetailGaze(val_images_dir, val_pickle_path, 'train') +val_set = RetailGaze2(val_images_dir, val_pickle_path, 'train') val_data_loader = DataLoader(dataset=val_set, batch_size=4, shuffle=True, num_workers=16) print ('Test') -test_set = RetailGaze(test_images_dir, test_pickle_path, 'test') +test_set = RetailGaze2(test_images_dir, test_pickle_path, 'test') test_data_loader = DataLoader(test_set, batch_size=batch_size//2, shuffle=False, num_workers=8) @@ -118,8 +118,8 @@ from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter('runs/face_3d') -# train_face3d(model, train_data_loader, val_data_loader, criterion, optimizer, logger, writer, num_epochs=50, patience=10) -test_face3d(model, test_data_loader, logger, test_depth=False) +train_face3d(model, train_data_loader, val_data_loader, criterion, optimizer, logger, writer, num_epochs=50, patience=10) +# test_face3d(model, test_data_loader, logger, test_depth=False) # train_face_depth(model_ft2, train_data_loader, val_data_loader, criterion, optimizer, logger, writer, num_epochs=50, patience=10) # test_face_depth(model_ft2, test_data_loader, logger) # train_face3d_bias(model_ft2, train_data_loader, val_data_loader, criterion, optimizer, logger, writer, num_epochs=50, patience=10)