custom_controlnet_aux/mesh_graphormer/pipeline.py

import os
import torch
import gc
import numpy as np
from custom_controlnet_aux.mesh_graphormer.depth_preprocessor import Preprocessor

import torchvision.models as models
from custom_mesh_graphormer.modeling.bert import BertConfig, Graphormer
from custom_mesh_graphormer.modeling.bert import Graphormer_Hand_Network as Graphormer_Network
from custom_mesh_graphormer.modeling._mano import MANO, Mesh
from custom_mesh_graphormer.modeling.hrnet.hrnet_cls_net_gridfeat import get_cls_net_gridfeat
from custom_mesh_graphormer.modeling.hrnet.config import config as hrnet_config
from custom_mesh_graphormer.modeling.hrnet.config import update_config as hrnet_update_config
from custom_mesh_graphormer.utils.miscellaneous import set_seed
from argparse import Namespace
from pathlib import Path
import cv2
from torchvision import transforms
import numpy as np
import cv2
from trimesh import Trimesh
from trimesh.ray.ray_triangle import RayMeshIntersector
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from torchvision import transforms
from pathlib import Path
from custom_controlnet_aux.util import custom_hf_download
import custom_mesh_graphormer
from comfy.model_management import soft_empty_cache
from packaging import version

args = Namespace(
    num_workers=4,
    img_scale_factor=1,
    image_file_or_path=os.path.join('', 'MeshGraphormer', 'samples', 'hand'), 
    model_name_or_path=str(Path(custom_mesh_graphormer.__file__).parent / "modeling/bert/bert-base-uncased"),
    resume_checkpoint=None,
    output_dir='output/',
    config_name='',
    a='hrnet-w64',
    arch='hrnet-w64',
    num_hidden_layers=4,
    hidden_size=-1,
    num_attention_heads=4,
    intermediate_size=-1,
    input_feat_dim='2051,512,128',
    hidden_feat_dim='1024,256,64',
    which_gcn='0,0,1',
    mesh_type='hand',
    run_eval_only=True,
    device="cpu",
    seed=88,
    hrnet_checkpoint=custom_hf_download("hr16/ControlNet-HandRefiner-pruned", 'hrnetv2_w64_imagenet_pretrained.pth')
)

#Since mediapipe v0.10.5, the hand category has been correct
if version.parse(mp.__version__) >= version.parse('0.10.5'):
    true_hand_category = {"Right": "right", "Left": "left"}
else:
    true_hand_category = {"Right": "left", "Left": "right"}

class MeshGraphormerMediapipe(Preprocessor):
    def __init__(self, args=args, detect_thr=0.6, presence_thr=0.6) -> None:
        #global logger
        # Setup CUDA, GPU & distributed training
        args.num_gpus = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
        os.environ['OMP_NUM_THREADS'] = str(args.num_workers)
        print('set os.environ[OMP_NUM_THREADS] to {}'.format(os.environ['OMP_NUM_THREADS']))

        #mkdir(args.output_dir)
        #logger = setup_logger("Graphormer", args.output_dir, get_rank())
        set_seed(args.seed, args.num_gpus)
        #logger.info("Using {} GPUs".format(args.num_gpus))

        # Mesh and MANO utils
        mano_model = MANO().to(args.device)
        mano_model.layer = mano_model.layer.to(args.device)
        mesh_sampler = Mesh(device=args.device)

        # Renderer for visualization
        # renderer = Renderer(faces=mano_model.face)

        # Load pretrained model
        trans_encoder = []

        input_feat_dim = [int(item) for item in args.input_feat_dim.split(',')]
        hidden_feat_dim = [int(item) for item in args.hidden_feat_dim.split(',')]
        output_feat_dim = input_feat_dim[1:] + [3]

        # which encoder block to have graph convs
        which_blk_graph = [int(item) for item in args.which_gcn.split(',')]

        if args.run_eval_only==True and args.resume_checkpoint!=None and args.resume_checkpoint!='None' and 'state_dict' not in args.resume_checkpoint:
            # if only run eval, load checkpoint
            #logger.info("Evaluation: Loading from checkpoint {}".format(args.resume_checkpoint))
            _model = torch.load(args.resume_checkpoint)

        else:
            # init three transformer-encoder blocks in a loop
            for i in range(len(output_feat_dim)):
                config_class, model_class = BertConfig, Graphormer
                config = config_class.from_pretrained(args.config_name if args.config_name \
                        else args.model_name_or_path)

                config.output_attentions = False
                config.img_feature_dim = input_feat_dim[i] 
                config.output_feature_dim = output_feat_dim[i]
                args.hidden_size = hidden_feat_dim[i]
                args.intermediate_size = int(args.hidden_size*2)

                if which_blk_graph[i]==1:
                    config.graph_conv = True
                    #logger.info("Add Graph Conv")
                else:
                    config.graph_conv = False

                config.mesh_type = args.mesh_type

                # update model structure if specified in arguments
                update_params = ['num_hidden_layers', 'hidden_size', 'num_attention_heads', 'intermediate_size']
                for idx, param in enumerate(update_params):
                    arg_param = getattr(args, param)
                    config_param = getattr(config, param)
                    if arg_param > 0 and arg_param != config_param:
                        #logger.info("Update config parameter {}: {} -> {}".format(param, config_param, arg_param))
                        setattr(config, param, arg_param)

                # init a transformer encoder and append it to a list
                assert config.hidden_size % config.num_attention_heads == 0
                model = model_class(config=config) 
                #logger.info("Init model from scratch.")
                trans_encoder.append(model)
            
            # create backbone model
            if args.arch=='hrnet':
                hrnet_yaml = Path(__file__).parent / 'cls_hrnet_w40_sgd_lr5e-2_wd1e-4_bs32_x100.yaml'
                hrnet_checkpoint = args.hrnet_checkpoint
                hrnet_update_config(hrnet_config, hrnet_yaml)
                backbone = get_cls_net_gridfeat(hrnet_config, pretrained=hrnet_checkpoint)
                #logger.info('=> loading hrnet-v2-w40 model')
            elif args.arch=='hrnet-w64':
                hrnet_yaml = Path(__file__).parent / 'cls_hrnet_w64_sgd_lr5e-2_wd1e-4_bs32_x100.yaml'
                hrnet_checkpoint = args.hrnet_checkpoint
                hrnet_update_config(hrnet_config, hrnet_yaml)
                backbone = get_cls_net_gridfeat(hrnet_config, pretrained=hrnet_checkpoint)
                #logger.info('=> loading hrnet-v2-w64 model')
            else:
                print("=> using pre-trained model '{}'".format(args.arch))
                backbone = models.__dict__[args.arch](pretrained=True)
                # remove the last fc layer
                backbone = torch.nn.Sequential(*list(backbone.children())[:-1])

            trans_encoder = torch.nn.Sequential(*trans_encoder)
            total_params = sum(p.numel() for p in trans_encoder.parameters())
            #logger.info('Graphormer encoders total parameters: {}'.format(total_params))
            backbone_total_params = sum(p.numel() for p in backbone.parameters())
            #logger.info('Backbone total parameters: {}'.format(backbone_total_params))

            # build end-to-end Graphormer network (CNN backbone + multi-layer Graphormer encoder)
            _model = Graphormer_Network(args, config, backbone, trans_encoder)

            if args.resume_checkpoint!=None and args.resume_checkpoint!='None':
                # for fine-tuning or resume training or inference, load weights from checkpoint
                #logger.info("Loading state dict from checkpoint {}".format(args.resume_checkpoint))
                # workaround approach to load sparse tensor in graph conv.
                state_dict = torch.load(args.resume_checkpoint)
                _model.load_state_dict(state_dict, strict=False)
                del state_dict
                gc.collect()
                soft_empty_cache()

        # update configs to enable attention outputs
        setattr(_model.trans_encoder[-1].config,'output_attentions', True)
        setattr(_model.trans_encoder[-1].config,'output_hidden_states', True)
        _model.trans_encoder[-1].bert.encoder.output_attentions = True
        _model.trans_encoder[-1].bert.encoder.output_hidden_states =  True
        for iter_layer in range(4):
            _model.trans_encoder[-1].bert.encoder.layer[iter_layer].attention.self.output_attentions = True
        for inter_block in range(3):
            setattr(_model.trans_encoder[-1].config,'device', args.device)

        _model.to(args.device)
        self._model = _model
        self.mano_model = mano_model
        self.mesh_sampler = mesh_sampler

        self.transform = transforms.Compose([           
                    transforms.ToTensor(),
                    transforms.Normalize(
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])])
        #Fix File loading is not yet supported on Windows
        with open(str( Path(__file__).parent / "hand_landmarker.task" ), 'rb') as file:
            model_data = file.read()
        base_options = python.BaseOptions(model_asset_buffer=model_data)
        options = vision.HandLandmarkerOptions(base_options=base_options,
                                            min_hand_detection_confidence=detect_thr,
                                            min_hand_presence_confidence=presence_thr,
                                            min_tracking_confidence=0.6,
                                            num_hands=2)

        self.detector = vision.HandLandmarker.create_from_options(options)
        
    
    def get_rays(self, W, H, fx, fy, cx, cy, c2w_t, center_pixels): # rot = I
   
        j, i = np.meshgrid(np.arange(H, dtype=np.float32), np.arange(W, dtype=np.float32))
        if center_pixels:
            i = i.copy() + 0.5
            j = j.copy() + 0.5

        directions = np.stack([(i - cx) / fx, (j - cy) / fy, np.ones_like(i)], -1)
        directions /= np.linalg.norm(directions, axis=-1, keepdims=True)

        rays_o = np.expand_dims(c2w_t,0).repeat(H*W, 0)

        rays_d = directions    # (H, W, 3)
        rays_d = (rays_d / np.linalg.norm(rays_d, axis=-1, keepdims=True)).reshape(-1,3)

        return rays_o, rays_d
    
    def get_mask_bounding_box(self, extrema, H, W, padding=30, dynamic_resize=0.15):
        x_min, x_max, y_min, y_max = extrema
        bb_xpad = max(int((x_max - x_min + 1) * dynamic_resize), padding)
        bb_ypad = max(int((y_max - y_min + 1) * dynamic_resize), padding)
        bbx_min = np.max((x_min - bb_xpad, 0))
        bbx_max = np.min((x_max + bb_xpad, W-1))
        bby_min = np.max((y_min - bb_ypad, 0))
        bby_max = np.min((y_max + bb_ypad, H-1))
        return bbx_min, bbx_max, bby_min, bby_max

    def run_inference(self, img, Graphormer_model, mano, mesh_sampler, scale, crop_len):
        global args
        H, W = int(crop_len), int(crop_len)
        Graphormer_model.eval()
        mano.eval()
        device = next(Graphormer_model.parameters()).device
        with torch.no_grad():
            img_tensor = self.transform(img)
            batch_imgs = torch.unsqueeze(img_tensor, 0).to(device)
            
            # forward-pass
            pred_camera, pred_3d_joints, pred_vertices_sub, pred_vertices, hidden_states, att = Graphormer_model(batch_imgs, mano, mesh_sampler)

            # obtain 3d joints, which are regressed from the full mesh
            pred_3d_joints_from_mesh = mano.get_3d_joints(pred_vertices)
            # obtain 2d joints, which are projected from 3d joints of mesh
            #pred_2d_joints_from_mesh = orthographic_projection(pred_3d_joints_from_mesh.contiguous(), pred_camera.contiguous())
            #pred_2d_coarse_vertices_from_mesh = orthographic_projection(pred_vertices_sub.contiguous(), pred_camera.contiguous())
            pred_camera = pred_camera.cpu()
            pred_vertices = pred_vertices.cpu()
            mesh = Trimesh(vertices=pred_vertices[0], faces=mano.face)
            res = crop_len
            focal_length = 1000 * scale
            camera_t = np.array([-pred_camera[1], -pred_camera[2], -2*focal_length/(res * pred_camera[0] +1e-9)])
            pred_3d_joints_camera = pred_3d_joints_from_mesh.cpu()[0] - camera_t
            z_3d_dist = pred_3d_joints_camera[:,2].clone()

            pred_2d_joints_img_space = ((pred_3d_joints_camera/z_3d_dist[:,None]) * np.array((focal_length, focal_length, 1)))[:,:2] + np.array((W/2, H/2))

            rays_o, rays_d = self.get_rays(W, H, focal_length, focal_length, W/2, H/2, camera_t, True)
            coords = np.array(list(np.ndindex(H,W))).reshape(H,W,-1).transpose(1,0,2).reshape(-1,2)
            intersector = RayMeshIntersector(mesh)
            points, index_ray, _ = intersector.intersects_location(rays_o, rays_d, multiple_hits=False)

            tri_index = intersector.intersects_first(rays_o, rays_d)

            tri_index = tri_index[index_ray]

            assert len(index_ray) == len(tri_index)
            
            discriminator = (np.sum(mesh.face_normals[tri_index]* rays_d[index_ray], axis=-1)<= 0)
            points = points[discriminator] # ray intesects in interior faces, discard them

            if len(points) == 0:
                return None, None
            depth = (points + camera_t)[:,-1]
            index_ray = index_ray[discriminator]
            pixel_ray = coords[index_ray]

            minval = np.min(depth)
            maxval = np.max(depth)
            depthmap = np.zeros([H,W])

            depthmap[pixel_ray[:, 0], pixel_ray[:, 1]] = 1.0 - (0.8 * (depth - minval) / (maxval - minval))
            depthmap *= 255
        return depthmap, pred_2d_joints_img_space


    def get_depth(self, np_image, padding):
        info = {}

        # STEP 3: Load the input image.
        #https://stackoverflow.com/a/76407270
        image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image.copy())

        # STEP 4: Detect hand landmarks from the input image.
        detection_result = self.detector.detect(image)

        handedness_list = detection_result.handedness
        hand_landmarks_list = detection_result.hand_landmarks

        raw_image = image.numpy_view()
        H, W, C = raw_image.shape


        # HANDLANDMARKS CAN BE EMPTY, HANDLE THIS!
        if len(hand_landmarks_list) == 0:
            return None, None, None
        raw_image = raw_image[:, :, :3]

        padded_image = np.zeros((H*2, W*2, 3))
        padded_image[int(1/2 * H):int(3/2 * H), int(1/2 * W):int(3/2 * W)] = raw_image

        hand_landmarks_list, handedness_list = zip(
            *sorted(
                zip(hand_landmarks_list, handedness_list), key=lambda x: x[0][9].z, reverse=True
            )
        )

        padded_depthmap = np.zeros((H*2, W*2))
        mask = np.zeros((H, W))
        crop_boxes = []
        #bboxes = []
        groundtruth_2d_keypoints = []
        hands = []
        depth_failure = False
        crop_lens = []
        abs_boxes = []
        
        for idx in range(len(hand_landmarks_list)):
            hand = true_hand_category[handedness_list[idx][0].category_name]
            hands.append(hand)
            hand_landmarks = hand_landmarks_list[idx]
            handedness = handedness_list[idx]
            height, width, _ = raw_image.shape
            x_coordinates = [landmark.x for landmark in hand_landmarks]
            y_coordinates = [landmark.y for landmark in hand_landmarks]

            # x_min, x_max, y_min, y_max: extrema from mediapipe keypoint detection
            x_min = int(min(x_coordinates) * width)
            x_max = int(max(x_coordinates) * width)
            x_c = (x_min + x_max)//2
            y_min = int(min(y_coordinates) * height)
            y_max = int(max(y_coordinates) * height)
            y_c = (y_min + y_max)//2
            abs_boxes.append([x_min, x_max, y_min, y_max])

            #if x_max - x_min < 60 or y_max - y_min < 60:
            #    continue

            crop_len = (max(x_max - x_min, y_max - y_min) * 1.6) //2 * 2

            # crop_x_min, crop_x_max, crop_y_min, crop_y_max: bounding box for mesh reconstruction 
            crop_x_min = int(x_c - (crop_len/2 - 1) + W/2)
            crop_x_max = int(x_c + crop_len/2 + W/2)
            crop_y_min = int(y_c - (crop_len/2 - 1) + H/2)
            crop_y_max = int(y_c + crop_len/2 + H/2)

            cropped = padded_image[crop_y_min:crop_y_max+1, crop_x_min:crop_x_max+1]
            crop_boxes.append([crop_y_min, crop_y_max, crop_x_min, crop_x_max])
            crop_lens.append(crop_len)
            if hand == "left":
                cropped = cv2.flip(cropped, 1)

            if crop_len < 224:
                graphormer_input = cv2.resize(cropped, (224, 224), interpolation=cv2.INTER_CUBIC)
            else:
                graphormer_input = cv2.resize(cropped, (224, 224), interpolation=cv2.INTER_AREA)
            scale = crop_len/224
            cropped_depthmap, pred_2d_keypoints = self.run_inference(graphormer_input.astype(np.uint8), self._model, self.mano_model, self.mesh_sampler, scale, int(crop_len)) 

            if cropped_depthmap is None:
                depth_failure = True
                break
            #keypoints_image_space = pred_2d_keypoints * (crop_y_max - crop_y_min + 1)/224
            groundtruth_2d_keypoints.append(pred_2d_keypoints)
            
            if hand == "left":
                cropped_depthmap = cv2.flip(cropped_depthmap, 1)
            resized_cropped_depthmap = cv2.resize(cropped_depthmap, (int(crop_len), int(crop_len)), interpolation=cv2.INTER_LINEAR)
            nonzero_y, nonzero_x = (resized_cropped_depthmap != 0).nonzero()
            if len(nonzero_y) == 0 or len(nonzero_x) == 0:
                depth_failure = True
                break
            padded_depthmap[crop_y_min+nonzero_y, crop_x_min+nonzero_x] = resized_cropped_depthmap[nonzero_y, nonzero_x]

            # nonzero stands for nonzero value on the depth map
            # coordinates of nonzero depth pixels in original image space
            original_nonzero_x = crop_x_min+nonzero_x - int(W/2)
            original_nonzero_y = crop_y_min+nonzero_y - int(H/2)
            
            nonzerox_min = min(np.min(original_nonzero_x), x_min)
            nonzerox_max = max(np.max(original_nonzero_x), x_max)
            nonzeroy_min = min(np.min(original_nonzero_y), y_min)
            nonzeroy_max = max(np.max(original_nonzero_y), y_max)

            bbx_min, bbx_max, bby_min, bby_max = self.get_mask_bounding_box((nonzerox_min, nonzerox_max, nonzeroy_min, nonzeroy_max), H, W, padding)
            mask[bby_min:bby_max+1, bbx_min:bbx_max+1] = 1.0
            #bboxes.append([int(bbx_min), int(bbx_max), int(bby_min), int(bby_max)])
        if depth_failure:
            #print("cannot detect normal hands")
            return None, None, None
        depthmap = padded_depthmap[int(1/2 * H):int(3/2 * H), int(1/2 * W):int(3/2 * W)].astype(np.uint8)
        mask = (255.0 * mask).astype(np.uint8)
        info["groundtruth_2d_keypoints"] = groundtruth_2d_keypoints
        info["hands"] = hands
        info["crop_boxes"] = crop_boxes
        info["crop_lens"] = crop_lens
        info["abs_boxes"] = abs_boxes
        return depthmap, mask, info
    
    def get_keypoints(self, img, Graphormer_model, mano, mesh_sampler, scale, crop_len):
        global args
        H, W = int(crop_len), int(crop_len)
        Graphormer_model.eval()
        mano.eval()
        device = next(Graphormer_model.parameters()).device
        with torch.no_grad():
            img_tensor = self.transform(img)
            #print(img_tensor)
            batch_imgs = torch.unsqueeze(img_tensor, 0).to(device)
            
            # forward-pass
            pred_camera, pred_3d_joints, pred_vertices_sub, pred_vertices, hidden_states, att = Graphormer_model(batch_imgs, mano, mesh_sampler)

            # obtain 3d joints, which are regressed from the full mesh
            pred_3d_joints_from_mesh = mano.get_3d_joints(pred_vertices)
            # obtain 2d joints, which are projected from 3d joints of mesh
            #pred_2d_joints_from_mesh = orthographic_projection(pred_3d_joints_from_mesh.contiguous(), pred_camera.contiguous())
            #pred_2d_coarse_vertices_from_mesh = orthographic_projection(pred_vertices_sub.contiguous(), pred_camera.contiguous())
            pred_camera = pred_camera.cpu()
            pred_vertices = pred_vertices.cpu()
            #
            res = crop_len
            focal_length = 1000 * scale
            camera_t = np.array([-pred_camera[1], -pred_camera[2], -2*focal_length/(res * pred_camera[0] +1e-9)])
            pred_3d_joints_camera = pred_3d_joints_from_mesh.cpu()[0] - camera_t
            z_3d_dist = pred_3d_joints_camera[:,2].clone()
            pred_2d_joints_img_space = ((pred_3d_joints_camera/z_3d_dist[:,None]) * np.array((focal_length, focal_length, 1)))[:,:2] + np.array((W/2, H/2))
            
        return pred_2d_joints_img_space
    

    def eval_mpjpe(self, sample, info):
        H, W, C = sample.shape
        padded_image = np.zeros((H*2, W*2, 3))
        padded_image[int(1/2 * H):int(3/2 * H), int(1/2 * W):int(3/2 * W)] = sample
        crop_boxes = info["crop_boxes"]
        hands = info["hands"]
        groundtruth_2d_keypoints = info["groundtruth_2d_keypoints"]
        crop_lens = info["crop_lens"]
        pjpe = 0
        for i in range(len(crop_boxes)):#box in crop_boxes:
            crop_y_min, crop_y_max, crop_x_min, crop_x_max = crop_boxes[i]
            cropped = padded_image[crop_y_min:crop_y_max+1, crop_x_min:crop_x_max+1]
            hand = hands[i]
            if hand == "left":
                cropped = cv2.flip(cropped, 1)
            crop_len = crop_lens[i]
            scale = crop_len/224
            if crop_len < 224:
                graphormer_input = cv2.resize(cropped, (224, 224), interpolation=cv2.INTER_CUBIC)
            else:
                graphormer_input = cv2.resize(cropped, (224, 224), interpolation=cv2.INTER_AREA)
            generated_keypoint = self.get_keypoints(graphormer_input.astype(np.uint8), self._model, self.mano_model, self.mesh_sampler, scale, crop_len)
            #generated_keypoint = generated_keypoint * ((crop_y_max - crop_y_min + 1)/224)
            pjpe += np.sum(np.sqrt(np.sum(((generated_keypoint - groundtruth_2d_keypoints[i]) ** 2).numpy(), axis=1)))
            pass
        mpjpe = pjpe/(len(crop_boxes) * 21)
        return mpjpe