face3d dataloader updated

PrimeshShamilka · Feb 4, 2022 · 28cdd12 · 28cdd12
1 parent bb577aa
commit 28cdd12
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git 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
@@ -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)