基于 PyTorch 的人臉關(guān)鍵點(diǎn)檢測

if not os.path.exists('/content/ibug_300W_large_face_landmark_dataset'):    !wget http://dlib.net/files/data/ibug_300W_large_face_landmark_dataset.tar.gz    !tar -xvzf 'ibug_300W_large_face_landmark_dataset.tar.gz'        !rm -r 'ibug_300W_large_face_landmark_dataset.tar.gz'

class Transforms():    def __init__(self):        pass        def rotate(self, image, landmarks, angle):        angle = random.uniform(-angle, +angle)
        transformation_matrix = torch.tensor([            [+cos(radians(angle)), -sin(radians(angle))],             [+sin(radians(angle)), +cos(radians(angle))]        ])
        image = imutils.rotate(np.array(image), angle)
        landmarks = landmarks - 0.5        new_landmarks = np.matmul(landmarks, transformation_matrix)        new_landmarks = new_landmarks + 0.5        return Image.fromarray(image), new_landmarks
    def resize(self, image, landmarks, img_size):        image = TF.resize(image, img_size)        return image, landmarks
    def color_jitter(self, image, landmarks):        color_jitter = transforms.ColorJitter(brightness=0.3,                                               contrast=0.3,                                              saturation=0.3,                                               hue=0.1)        image = color_jitter(image)        return image, landmarks
    def crop_face(self, image, landmarks, crops):        left = int(crops['left'])        top = int(crops['top'])        width = int(crops['width'])        height = int(crops['height'])
        image = TF.crop(image, top, left, height, width)
        img_shape = np.array(image).shape        landmarks = torch.tensor(landmarks) - torch.tensor([[left, top]])        landmarks = landmarks / torch.tensor([img_shape[1], img_shape[0]])        return image, landmarks
    def __call__(self, image, landmarks, crops):        image = Image.fromarray(image)        image, landmarks = self.crop_face(image, landmarks, crops)        image, landmarks = self.resize(image, landmarks, (224, 224))        image, landmarks = self.color_jitter(image, landmarks)        image, landmarks = self.rotate(image, landmarks, angle=10)                image = TF.to_tensor(image)        image = TF.normalize(image, [0.5], [0.5])        return image, landmarks

class FaceLandmarksDataset(Dataset):
    def __init__(self, transform=None):
        tree = ET.parse('ibug_300W_large_face_landmark_dataset/labels_ibug_300W_train.xml')        root = tree.getroot()
        self.image_filenames = []        self.landmarks = []        self.crops = []        self.transform = transform        self.root_dir = 'ibug_300W_large_face_landmark_dataset'                for filename in root[2]:            self.image_filenames.append(os.path.join(self.root_dir, filename.attrib['file']))
            self.crops.append(filename[0].attrib)
            landmark = []            for num in range(68):                x_coordinate = int(filename[0][num].attrib['x'])                y_coordinate = int(filename[0][num].attrib['y'])                landmark.append([x_coordinate, y_coordinate])            self.landmarks.append(landmark)
        self.landmarks = np.array(self.landmarks).astype('float32')     
        assert len(self.image_filenames) == len(self.landmarks)
    def __len__(self):        return len(self.image_filenames)
    def __getitem__(self, index):        image = cv2.imread(self.image_filenames[index], 0)        landmarks = self.landmarks[index]                if self.transform:            image, landmarks = self.transform(image, landmarks, self.crops[index])
        landmarks = landmarks - 0.5
        return image, landmarks
dataset = FaceLandmarksDataset(Transforms())

class Network(nn.Module):    def __init__(self,num_classes=136):        super().__init__()        self.model_name='resnet18'        self.model=models.resnet18()        self.model.conv1=nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)        self.model.fc=nn.Linear(self.model.fc.in_features, num_classes)            def forward(self, x):        x=self.model(x)        return x

network = Network()network.cuda()    
criterion = nn.MSELoss()optimizer = optim.Adam(network.parameters(), lr=0.0001)
loss_min = np.infnum_epochs = 10
start_time = time.time()for epoch in range(1,num_epochs+1):        loss_train = 0    loss_valid = 0    running_loss = 0        network.train()    for step in range(1,len(train_loader)+1):            images, landmarks = next(iter(train_loader))                images = images.cuda()        landmarks = landmarks.view(landmarks.size(0),-1).cuda()                 predictions = network(images)                # clear all the gradients before calculating them        optimizer.zero_grad()                # find the loss for the current step        loss_train_step = criterion(predictions, landmarks)                # calculate the gradients        loss_train_step.backward()                # update the parameters        optimizer.step()                loss_train += loss_train_step.item()        running_loss = loss_train/step                print_overwrite(step, len(train_loader), running_loss, 'train')            network.eval()     with torch.no_grad():                for step in range(1,len(valid_loader)+1):                        images, landmarks = next(iter(valid_loader))                    images = images.cuda()            landmarks = landmarks.view(landmarks.size(0),-1).cuda()                    predictions = network(images)
            # find the loss for the current step            loss_valid_step = criterion(predictions, landmarks)
            loss_valid += loss_valid_step.item()            running_loss = loss_valid/step
            print_overwrite(step, len(valid_loader), running_loss, 'valid')        loss_train /= len(train_loader)    loss_valid /= len(valid_loader)        print('\n--------------------------------------------------')    print('Epoch: {}  Train Loss: {:.4f}  Valid Loss: {:.4f}'.format(epoch, loss_train, loss_valid))    print('--------------------------------------------------')        if loss_valid < loss_min:        loss_min = loss_valid        torch.save(network.state_dict(), '/content/face_landmarks.pth')         print("\nMinimum Validation Loss of {:.4f} at epoch {}/{}".format(loss_min, epoch, num_epochs))        print('Model Saved\n')     print('Training Complete')print("Total Elapsed Time : {} s".format(time.time()-start_time))

import timeimport cv2import osimport numpy as npimport matplotlib.pyplot as pltfrom PIL import Imageimport imutils
import torchimport torch.nn as nnfrom torchvision import modelsimport torchvision.transforms.functional as TF#######################################################################image_path = 'pic.jpg'weights_path = 'face_landmarks.pth'frontal_face_cascade_path = 'haarcascade_frontalface_default.xml'#######################################################################class Network(nn.Module):    def __init__(self,num_classes=136):        super().__init__()        self.model_name='resnet18'        self.model=models.resnet18(pretrained=False)        self.model.conv1=nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)        self.model.fc=nn.Linear(self.model.fc.in_features,num_classes)            def forward(self, x):        x=self.model(x)        return x
#######################################################################face_cascade = cv2.CascadeClassifier(frontal_face_cascade_path)
best_network = Network()best_network.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu'))) best_network.eval()
image = cv2.imread(image_path)grayscale_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)display_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)height, width,_ = image.shape
faces = face_cascade.detectMultiScale(grayscale_image, 1.1, 4)
all_landmarks = []for (x, y, w, h) in faces:    image = grayscale_image[y:y+h, x:x+w]    image = TF.resize(Image.fromarray(image), size=(224, 224))    image = TF.to_tensor(image)    image = TF.normalize(image, [0.5], [0.5])
    with torch.no_grad():        landmarks = best_network(image.unsqueeze(0)) 
    landmarks = (landmarks.view(68,2).detach().numpy() + 0.5) * np.array([[w, h]]) + np.array([[x, y]])    all_landmarks.append(landmarks)
plt.figure()plt.imshow(display_image)for landmarks in all_landmarks:    plt.scatter(landmarks[:,0], landmarks[:,1], c = 'c', s = 5)
plt.show()