georgefen/Face-Landmark-ControlNet

1import cv22import numpy as np3import torch4import os5 6from einops import rearrange7from .models.mbv2_mlsd_tiny import MobileV2_MLSD_Tiny8from .models.mbv2_mlsd_large import MobileV2_MLSD_Large9from .utils import pred_lines10 11from annotator.util import annotator_ckpts_path12 13 14remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_large_512_fp32.pth"15 16 17class MLSDdetector:18    def __init__(self):19        model_path = os.path.join(annotator_ckpts_path, "mlsd_large_512_fp32.pth")20        if not os.path.exists(model_path):21            from basicsr.utils.download_util import load_file_from_url22            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)23        model = MobileV2_MLSD_Large()24        model.load_state_dict(torch.load(model_path), strict=True)25        self.model = model.cuda().eval()26 27    def __call__(self, input_image, thr_v, thr_d):28        assert input_image.ndim == 329        img = input_image30        img_output = np.zeros_like(img)31        try:32            with torch.no_grad():33                lines = pred_lines(img, self.model, [img.shape[0], img.shape[1]], thr_v, thr_d)34                for line in lines:35                    x_start, y_start, x_end, y_end = [int(val) for val in line]36                    cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1)37        except Exception as e:38            pass39        return img_output[:, :, 0]40