flame_tracking_single_image.py

import argparse
import json
import os
import time
from pathlib import Path

import cv2
import numpy as np
import torch
import torchvision
import tyro
import yaml
from loguru import logger
from PIL import Image

from external.human_matting import StyleMatteEngine as HumanMattingEngine
from external.landmark_detection.FaceBoxesV2.faceboxes_detector import \
    FaceBoxesDetector
from external.landmark_detection.infer_image import Alignment
from external.vgghead_detector import VGGHeadDetector
from vhap.config.base import BaseTrackingConfig
from vhap.export_as_nerf_dataset import (NeRFDatasetWriter,
                                         TrackedFLAMEDatasetWriter, split_json)
from vhap.model.tracker import GlobalTracker

# Define error codes for various processing failures.
ERROR_CODE = {'FailedToDetect': 1, 'FailedToOptimize': 2, 'FailedToExport': 3}


def expand_bbox(bbox, scale=1.1):
    """Expands the bounding box by a given scale."""
    xmin, ymin, xmax, ymax = bbox.unbind(dim=-1)
    center_x, center_y = (xmin + xmax) / 2, (ymin + ymax) / 2
    extension_size = torch.sqrt((ymax - ymin) * (xmax - xmin)) * scale
    x_min_expanded = center_x - extension_size / 2
    x_max_expanded = center_x + extension_size / 2
    y_min_expanded = center_y - extension_size / 2
    y_max_expanded = center_y + extension_size / 2
    return torch.stack(
        [x_min_expanded, y_min_expanded, x_max_expanded, y_max_expanded],
        dim=-1)


def load_config(src_folder: Path):
    """Load configuration from the given source folder."""
    config_file_path = src_folder / 'config.yml'
    if not config_file_path.exists():
        src_folder = sorted(
            src_folder.iterdir())[-1]  # Get the last modified folder
        config_file_path = src_folder / 'config.yml'
    assert config_file_path.exists(), f'File not found: {config_file_path}'

    config_data = yaml.load(config_file_path.read_text(), Loader=yaml.Loader)
    return src_folder, config_data


class FlameTrackingSingleImage:
    """Class for tracking and processing a single image."""
    def __init__(
            self,
            output_dir,
            alignment_model_path='./pretrain_model/68_keypoints_model.pkl',
            vgghead_model_path='./pretrain_model/vgghead/vgg_heads_l.trcd',
            human_matting_path='./pretrain_model/matting/stylematte_synth.pt',
            facebox_model_path='./pretrain_model/FaceBoxesV2.pth',
            detect_iris_landmarks=False):

        logger.info(f'Output Directory: {output_dir}')

        start_time = time.time()
        logger.info('Loading Pre-trained Models...')

        self.output_dir = output_dir
        self.output_preprocess = os.path.join(output_dir, 'preprocess')
        self.output_tracking = os.path.join(output_dir, 'tracking')
        self.output_export = os.path.join(output_dir, 'export')
        self.device = 'cuda:0'

        # Load alignment model
        assert os.path.exists(
            alignment_model_path), f'{alignment_model_path} does not exist!'
        args = self._parse_args()
        args.model_path = alignment_model_path
        self.alignment = Alignment(args,
                                   alignment_model_path,
                                   dl_framework='pytorch',
                                   device_ids=[0])

        # Load VGG head model
        assert os.path.exists(
            vgghead_model_path), f'{vgghead_model_path} does not exist!'
        self.vgghead_encoder = VGGHeadDetector(
            device=self.device, vggheadmodel_path=vgghead_model_path)

        # Load human matting model
        assert os.path.exists(
            human_matting_path), f'{human_matting_path} does not exist!'
        self.matting_engine = HumanMattingEngine(
            device=self.device, human_matting_path=human_matting_path)

        # Load face box detector model
        assert os.path.exists(
            facebox_model_path), f'{facebox_model_path} does not exist!'
        self.detector = FaceBoxesDetector('FaceBoxes', facebox_model_path,
                                          True, self.device)

        self.detect_iris_landmarks_flag = detect_iris_landmarks
        if self.detect_iris_landmarks_flag:
            from fdlite import FaceDetection, FaceLandmark, IrisLandmark
            self.iris_detect_faces = FaceDetection()
            self.iris_detect_face_landmarks = FaceLandmark()
            self.iris_detect_iris_landmarks = IrisLandmark()

        end_time = time.time()
        torch.cuda.empty_cache()
        logger.info(f'Finished Loading Pre-trained Models. Time: '
                    f'{end_time - start_time:.2f}s')

    def _parse_args(self):
        parser = argparse.ArgumentParser(description='Evaluation script')
        parser.add_argument('--output_dir',
                            type=str,
                            help='Output directory',
                            default='output')
        parser.add_argument('--config_name',
                            type=str,
                            help='Configuration name',
                            default='alignment')
        return parser.parse_args()

    def preprocess(self, input_image_path):
        """Preprocess the input image for tracking."""
        if not os.path.exists(input_image_path):
            logger.warning(f'{input_image_path} does not exist!')
            return ERROR_CODE['FailedToDetect']

        start_time = time.time()
        logger.info('Starting Preprocessing...')
        name_list = []
        frame_index = 0

        # Bounding box detection
        frame = torchvision.io.read_image(input_image_path)
        try:
            _, frame_bbox, _ = self.vgghead_encoder(frame, frame_index)
        except Exception:
            logger.error('Failed to detect face')
            return ERROR_CODE['FailedToDetect']

        if frame_bbox is None:
            logger.error('Failed to detect face')
            return ERROR_CODE['FailedToDetect']

        # Expand bounding box
        name_list.append('00000.png')
        frame_bbox = expand_bbox(frame_bbox, scale=1.65).long()

        # Crop and resize
        cropped_frame = torchvision.transforms.functional.crop(
            frame,
            top=frame_bbox[1],
            left=frame_bbox[0],
            height=frame_bbox[3] - frame_bbox[1],
            width=frame_bbox[2] - frame_bbox[0])
        cropped_frame = torchvision.transforms.functional.resize(
            cropped_frame, (1024, 1024), antialias=True)

        # Apply matting
        cropped_frame, mask = self.matting_engine(cropped_frame / 255.0,
                                                  return_type='matting',
                                                  background_rgb=1.0)
        cropped_frame = cropped_frame.cpu() * 255.0
        saved_image = np.round(cropped_frame.cpu().permute(
            1, 2, 0).numpy()).astype(np.uint8)[:, :, (2, 1, 0)]

        # Create output directories if not exist
        self.sub_output_dir = os.path.join(
            self.output_preprocess,
            os.path.splitext(os.path.basename(input_image_path))[0])
        output_image_dir = os.path.join(self.sub_output_dir, 'images')
        output_mask_dir = os.path.join(self.sub_output_dir, 'mask')
        output_alpha_map_dir = os.path.join(self.sub_output_dir, 'alpha_maps')

        os.makedirs(output_image_dir, exist_ok=True)
        os.makedirs(output_mask_dir, exist_ok=True)
        os.makedirs(output_alpha_map_dir, exist_ok=True)

        # Save processed image, mask and alpha map
        cv2.imwrite(os.path.join(output_image_dir, name_list[frame_index]),
                    saved_image)
        cv2.imwrite(os.path.join(output_mask_dir, name_list[frame_index]),
                    np.array((mask.cpu() * 255.0)).astype(np.uint8))
        cv2.imwrite(
            os.path.join(output_alpha_map_dir,
                         name_list[frame_index]).replace('.png', '.jpg'),
            (np.ones_like(saved_image) * 255).astype(np.uint8))

        # Landmark detection
        detections, _ = self.detector.detect(saved_image, 0.8, 1)
        for idx, detection in enumerate(detections):
            x1_ori, y1_ori = detection[2], detection[3]
            x2_ori, y2_ori = x1_ori + detection[4], y1_ori + detection[5]

            scale = max(x2_ori - x1_ori, y2_ori - y1_ori) / 180
            center_w, center_h = (x1_ori + x2_ori) / 2, (y1_ori + y2_ori) / 2
            scale, center_w, center_h = float(scale), float(center_w), float(
                center_h)

            face_landmarks = self.alignment.analyze(saved_image, scale,
                                                    center_w, center_h)

        # Normalize and save landmarks
        normalized_landmarks = np.zeros((face_landmarks.shape[0], 3))
        normalized_landmarks[:, :2] = face_landmarks / 1024

        landmark_output_dir = os.path.join(self.sub_output_dir, 'landmark2d')
        os.makedirs(landmark_output_dir, exist_ok=True)

        landmark_data = {
            'bounding_box': [],
            'face_landmark_2d': normalized_landmarks[None, ...],
        }

        landmark_path = os.path.join(landmark_output_dir, 'landmarks.npz')
        np.savez(landmark_path, **landmark_data)

        if self.detect_iris_landmarks_flag:
            self._detect_iris_landmarks(
                os.path.join(output_image_dir, name_list[frame_index]))

        end_time = time.time()
        torch.cuda.empty_cache()
        logger.info(
            f'Finished Processing Image. Time: {end_time - start_time:.2f}s')

        return 0

    def optimize(self):
        """Optimize the tracking model using configuration data."""
        start_time = time.time()
        logger.info('Starting Optimization...')

        tyro.extras.set_accent_color('bright_yellow')
        config_data = tyro.cli(BaseTrackingConfig)

        config_data.data.sequence = self.sub_output_dir.split('/')[-1]
        config_data.data.root_folder = Path(
            os.path.dirname(self.sub_output_dir))

        if not os.path.exists(self.sub_output_dir):
            logger.error(f'Failed to load {self.sub_output_dir}')
            return ERROR_CODE['FailedToOptimize']

        config_data.exp.output_folder = Path(self.output_tracking)
        tracker = GlobalTracker(config_data)
        tracker.optimize()

        end_time = time.time()
        torch.cuda.empty_cache()
        logger.info(
            f'Finished Optimization. Time: {end_time - start_time:.2f}s')

        return 0

    def _detect_iris_landmarks(self, image_path):
        """Detect iris landmarks in the given image."""
        from fdlite import face_detection_to_roi, iris_roi_from_face_landmarks

        img = Image.open(image_path)
        img_size = (1024, 1024)

        face_detections = self.iris_detect_faces(img)
        if len(face_detections) != 1:
            logger.warning('Empty iris landmarks')
        else:
            face_detection = face_detections[0]
            try:
                face_roi = face_detection_to_roi(face_detection, img_size)
            except ValueError:
                logger.warning('Empty iris landmarks')
                return

            face_landmarks = self.iris_detect_face_landmarks(img, face_roi)
            if len(face_landmarks) == 0:
                logger.warning('Empty iris landmarks')
                return

            iris_rois = iris_roi_from_face_landmarks(face_landmarks, img_size)

            if len(iris_rois) != 2:
                logger.warning('Empty iris landmarks')
                return

            landmarks = []
            for iris_roi in iris_rois[::-1]:
                try:
                    iris_landmarks = self.iris_detect_iris_landmarks(
                        img, iris_roi).iris[0:1]
                except np.linalg.LinAlgError:
                    logger.warning('Failed to get iris landmarks')
                    break

                # For each landmark, append x and y coordinates scaled to 1024.
                for landmark in iris_landmarks:
                    landmarks.append(landmark.x * 1024)
                    landmarks.append(landmark.y * 1024)

            landmark_data = {'00000.png': landmarks}
            json.dump(
                landmark_data,
                open(
                    os.path.join(self.sub_output_dir, 'landmark2d',
                                 'iris.json'), 'w'))

    def export(self):
        """Export the tracking results to configured folder."""
        logger.info(f'Beginning export from {self.output_tracking}')
        start_time = time.time()
        if not os.path.exists(self.output_tracking):
            logger.error(f'Failed to load {self.output_tracking}')
            return ERROR_CODE['FailedToExport'], 'Failed'

        src_folder = Path(self.output_tracking)
        tgt_folder = Path(self.output_export,
                          self.sub_output_dir.split('/')[-1])
        src_folder, config_data = load_config(src_folder)

        nerf_writer = NeRFDatasetWriter(config_data.data, tgt_folder, None,
                                        None, 'white')
        nerf_writer.write()

        flame_writer = TrackedFLAMEDatasetWriter(config_data.model,
                                                 src_folder,
                                                 tgt_folder,
                                                 mode='param',
                                                 epoch=-1)
        flame_writer.write()

        split_json(tgt_folder)

        end_time = time.time()
        torch.cuda.empty_cache()
        logger.info(f'Finished Export. Time: {end_time - start_time:.2f}s')

        return 0, str(tgt_folder)