AI Project

checkpoint.py

@@ -23,13 +23,6 @@ def __init__(self, save_path):
             os.makedirs(self.save_path)
 
     def load_state(self, model, state_dict):
-        """
-        load state_dict to model
-        :params model:
-        :params state_dict:
-        :return: model
-        """
-        # set the model in evaluation mode, otherwise the accuracy will change
         model.eval()
         model_dict = model.state_dict()
 
@@ -49,7 +42,6 @@ def load_model(self, model_path):
         """
         if os.path.isfile(model_path):
             print("|===>Load retrain model from:", model_path)
-            # model_state_dict = torch.load(model_path, map_location={'cuda:1':'cuda:0'})
             model_state_dict = torch.load(model_path, map_location='cpu')
             return model_state_dict
         else:

data/test/crop_eye_with_landmarks.py

@@ -48,7 +48,7 @@ def crop_eye_with_landmarks(img, landmarks):
     large_rect = larger_left_eye_rect
     left_eye = img[large_rect[1]:large_rect[1]+large_rect[3],
                large_rect[0]:large_rect[0]+large_rect[2]]
-    larger_rect = larger_right_eye_rect
+    large_rect = larger_right_eye_rect
     right_eye = img[large_rect[1]:large_rect[1]+large_rect[3],
                large_rect[0]:large_rect[0]+large_rect[2]]
 

loss.py

@@ -3,18 +3,17 @@
 class NLL_OHEM(torch.nn.NLLLoss):
     """Online hard sample mining, Needs input from nn.LogSoftmax()"""
     def __init__(self, ratio):
-        super(NLL_OHEM, self).__init__(None, True):
-            self.ration = ratio
+        super(NLL_OHEM, self).__init__(None, True).self.ration = ratio
 
     def forward(self, x, y, ratio=None):
         if ratio is not None:
             self.ratio = ratio
         num_inst = x.size(0)
         num_hns = int(self.ratio*num_inst)
-        x_ = x.clone()
+        x = x.clone()
         inst_losses = torch.autograd.Variable(torch.zeros(num_inst)).cuda()
         for idx, label in enumerate(y.data):
-            insta_losses[idx] = -x.data[idx, label]
+            num_inst[idx] = -x.data[idx, label]
         _, idxs = inst_losses.topk(num_hns)
         y_hn = y.index_select(0, idxs)
         x_hn = y.index_select(0, idxs)

test/test_image.py

@@ -1,23 +1,27 @@
-import os, sys
-sys.path.append('.')
+import os
+import sys
 import time
 import pathlib
-
 import cv2
 import numpy as np
 import torch
-
-import config
 from nets.optimized_landmark import LNet
-from tools.utils import show_landmarks, show_tensor_landmarks, draw_tensor_landmarks
+from tools.utils import show_landmarks, show_tensor_landmarks
 
 if __name__ == "__main__":
     model = LNet()
-    # model.load_state_dict(torch.load('result/check_point/{0}_landmarks_model_200.pth'.format(config.NUM_LANDMARKS)))
+    pretrained_weights_path = './pretrained_weights/{0}_landmarks_model_200.pth'.format(config.NUM_LANDMARKS)
+
+    if os.path.exists(pretrained_weights_path):
+        model.load_state_dict(torch.load(pretrained_weights_path))
+        print("Loaded pretrained weights successfully.")
+    else:
+        print("Pretrained weights not found. Please check the path:", pretrained_weights_path)
+        sys.exit(1)
+
     model.eval()
-    model.load_state_dict(torch.load('./pretrained_weights/{0}_landmarks_model_200.pth'.format(config.NUM_LANDMARKS)))
+
     input_path = pathlib.Path('./data/test/')
-
     project_root = pathlib.Path()
     output_path = project_root / "result" / "imgs"
     output_path.mkdir(exist_ok=True)
@@ -27,22 +31,25 @@
         img = cv2.imread(str(input_img_filename))
         RGB_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         RGB_img = cv2.resize(RGB_img, (config.IMAGE_WIDTH, config.IMAGE_HEIGHT))
-        
+
         data = RGB_img - 127.5
         data = data / 127.5
-        
+
         data = data.transpose((2, 0, 1))
         data = np.expand_dims(data, axis=0)
         data = torch.Tensor(data)
-        with torch.no_grad(): 
+
+        with torch.no_grad():
             landmarks = model(data)
-        landmarks = landmarks.cpu().detach()
+
+        landmarks = landmarks.cpu().detach().numpy()[0]
 
-        show_tensor_landmarks(data[0], landmarks[0])
-
-        landmarks = [(i) for i in landmarks[0]]
-        h, w = img.shape[0:2]
-        landmarks = [(landmarks[2*i]*w, landmarks[2*i+1]*h) for i in range(len(landmarks)//2)]
-        show_landmarks(img, landmarks) 
+        # Visualize landmarks using the show_tensor_landmarks function
+        show_tensor_landmarks(data[0], landmarks)
 
+        # Convert normalized landmarks to image coordinates
+        h, w = img.shape[0:2]
+        landmarks_img_coords = [(int(landmarks[2*i] * w), int(landmarks[2*i+1] * h)) for i in range(len(landmarks)//2)]
 
+        # Visualize landmarks using the show_landmarks function
+        show_landmarks(img, landmarks_img_coords)

tools/iris_dataset.py

@@ -60,13 +60,15 @@ def __call__(self, sample):
 class RandomFlip(object):
     def __init__(self, prob):
         self.prob = prob
+
     def __call__(self, sample):
         image = sample['image']
         landmarks = sample['landmarks']
+
         if random.random() < self.prob:
-            image = cv2.flip(image, 1) # 1 for flip around y axis, 0 for x axis, -1 for both
-            # landmarks[:,0] = image.shape[1] - landmarks[:,0] # flip x coordinates
-            landmarks[:, 0] = 1 - landmarks[:, 0]
+            image = cv2.flip(image, 1)  # 1 for flip around y axis, 0 for x axis, -1 for both
+            landmarks[:, 0] = 1 - landmarks[:, 0]  # flip x coordinates
+
         return {'image': image, 'landmarks': landmarks}
 
 class Rescale(object):
@@ -79,9 +81,14 @@ def __call__(self, sample):
         ratio = np.random.randint(self.low, self.high)
         image = sample['image']
         landmarks = sample['landmarks']
-        h, w = image.shape[0:2]
-        image = cv2.resize(image, (int(w/ratio), int(h/ratio)))
-        image = cv2.resize(image, (w, h))
+        h, w = image.shape[:2]
+
+        # Downscale
+        small_image = cv2.resize(image, (int(w / ratio), int(h / ratio)))
+
+        # Upscale
+        image = cv2.resize(small_image, (w, h))
+
         return {'image': image, 'landmarks': landmarks}
 
 class RandomGaussianBlur(object):
@@ -195,6 +202,104 @@ def kps_to_landmarks(self, kps):
         landmarks = np.array(landmarks)
         return landmarks
 
+
+
+def unnormalize_landmark(landmarks, image):
+    h, w = image.shape[:2]
+    landmarks[:, 0] = landmarks[:, 0] * w
+    landmarks[:, 1] = landmarks[:, 1] * h
+    return landmarks
+
+def normalize_landmark(landmarks, image):
+    h, w = image.shape[:2]
+    landmarks[:, 0] = landmarks[:, 0] / w
+    landmarks[:, 1] = landmarks[:, 1] / h
+    return landmarks
+
+class RandomCropResize(object):
+    def __init__(self, output_size, resize_ratio):
+        assert isinstance(output_size, (int, tuple))
+        if isinstance(output_size, int):
+            self.output_size = (output_size, output_size)
+        else:
+            assert len(output_size) == 2
+            self.output_size = output_size
+        self.resize_ratio = resize_ratio
+
+    def __call__(self, sample):
+        image, landmarks = sample['image'], sample['landmarks']
+
+        resize = np.random.random()
+
+        h, w = image.shape[:2]
+        new_w, new_h = self.output_size
+
+        if resize < self.resize_ratio:
+            top = np.random.randint(0, h - new_h)
+            left = np.random.randint(0, w - new_w)
+
+            landmarks = unnormalize_landmark(landmarks, image)
+            image = image[top:top + new_h, left:left + new_w]
+            landmarks = landmarks - [left, top]
+            landmarks = normalize_landmark(landmarks, image)
+        else:
+            image = cv2.resize(image, (new_w, new_h))
+            # landmarks = landmarks * [new_w / w, new_h / h]
+
+        return {'image': image, "landmarks": landmarks}
+
+class RandomRotate(object):
+    def __init__(self, degree):
+        self.degree = degree
+
+    def __call__(self, sample):
+        image = sample['image']
+        landmarks = sample['landmarks']
+        h, w = image.shape[:2]
+        img_h, img_w = image.shape[:2]
+        center = (img_w // 2, img_h // 2)
+        random_degree = np.random.uniform(-self.degree, self.degree)
+        rot_mat = cv2.getRotationMatrix2D(center, random_degree, 1)
+        image_rotated = cv2.warpAffine(image, rot_mat, (img_w, img_h))
+
+        landmark_rotated = np.asarray([(rot_mat[0][0] * x * w + rot_mat[0][1] * y * h + rot_mat[0][2],
+                                        rot_mat[1][0] * x * w + rot_mat[1][1] * y * h + rot_mat[1][2])
+                                       for (x, y) in landmarks])
+
+        for i in range(landmark_rotated.shape[0] // 2):
+            landmark_rotated[2 * i] /= w
+            landmark_rotated[2 * i + 1] /= h
+
+        return {'image': image_rotated, "landmarks": landmark_rotated}
+
+class RandomMotionBlur(object):
+    def __init__(self, radius):
+        self.radius = radius
+        self.seq = iaa.Sequential([
+            iaa.Sometimes(0.2,
+                          iaa.MotionBlur(k=self.radius)
+                          )
+        ])
+
+    def __call__(self, sample):
+        image = sample['image']
+        landmarks = sample['landmarks']
+        landmarks = unnormalize_landmark(landmarks, image)
+        kps = self.landmarks_to_kps(image, landmarks)
+        img_aug, kps_aug = self.seq(image=image, keypoints=kps)
+        landmarks_aug = self.kps_to_landmarks(kps_aug)
+        landmarks_aug = normalize_landmark(landmarks_aug, img_aug)
+        return {'image': img_aug, 'landmarks': landmarks_aug}
+
+    def landmarks_to_kps(self, image, landmarks):
+        kp_list = [Keypoint(x=landmarks[i][0], y=landmarks[i][1]) for i in range(landmarks.shape[0])]
+        kps = KeypointsOnImage(kp_list, shape=image.shape)
+        return kps
+
+    def kps_to_landmarks(self, kps):
+        landmarks = np.array([(kp.x_int, kp.y_int) for kp in kps.keypoints])
+        return landmarks
+
 class ToTensor(object):
     # def __init__(self, image_size):
     #    self.image_size = image_size