change_hair/project/hair_service_sd/hair_grow_swap.py

# -*- coding: utf-8 -*-
"""生发模块（走换发型工作流）

和换发型 /api/swapHair/v1 完全一致的工作流：infer_hairstyle_diy_jy 粗推理
→ warpAffine → photo_service + LoRA → webui inpainting → 贴回原图。
唯一区别：mask 是 origin_matting ∪ new_matting ∪ 手绘mask 的并集，
不减刘海/眼眶。

核心函数:
    hair_grow_swap(origin_img, hand_mask, hair_id, hairstyle_process,
                   landmark_processor, task_id, is_hr) -> BGR ndarray
"""
import os
import json
import glob
import time
import shutil

import cv2
import numpy as np

# 复用换发型的工具链
from gen_super_image import webui_img2img
from common.logger import config
from utils.landmark_processor import high_quality_warpAffine


def hair_grow_swap(origin_img, hand_mask, hair_id, hairstyle_process,
                   landmark_processor, task_id, is_hr=True):
    """在用户涂抹的缺发区生发（走换发型工作流 + 该发型 LoRA）。

    参数:
        origin_img: 用户原图 BGR ndarray
        hand_mask: 用户手绘mask，灰度 ndarray，白色(255)=生发区，与 origin_img 同尺寸
        hair_id: 选择的发型ID（用于加载 LoRA + 读材质/prompt）
        hairstyle_process: hair_init 初始化的 HairstyleModel 实例
        landmark_processor: landmark_processor 模块（含 high_quality_warpAffine）
        task_id: 任务ID（唯一字符串）
        is_hr: 是否高清模式（True=1152x1536，False=576x768）

    返回:
        BGR ndarray 生发结果图（与原图同尺寸）
    """
    start_all = time.time()

    # ===== 路径配置（与 swapHair 一致）=====
    hairstyle_dir = config.get('default', 'hairstyleDir')          # ref_hairstyle/
    user_dir = config.get('default', 'userDir')
    train_dir = config.get('default', 'train_dir')                 # train_material/
    userInfo_dir = config.get('default', 'userInfo_dir')
    res_dir = config.get('default', 'res_dir')

    hair_material_dir = os.path.join(train_dir, hair_id)           # train_material/<id>
    material_save_path = os.path.join(hairstyle_dir, hair_id)      # ref_hairstyle/<id>

    # 校验：ref 材质必须存在（infer_hairstyle_diy_jy 依赖）
    ref_img_path = os.path.join(material_save_path, "ref_rgb_8uc3_768.png")
    if not os.path.exists(ref_img_path):
        raise FileNotFoundError(f"发型材质不存在: {ref_img_path}")

    # 保存用户原图到工作目录（infer_hairstyle_diy_jy 需要读文件路径）
    user_img_name = f"hairgrow_{task_id}.jpg"
    new_user_img_path = os.path.join(user_dir, user_img_name)
    if os.path.exists(new_user_img_path):
        os.remove(new_user_img_path)
    cv2.imwrite(new_user_img_path, origin_img)
    origin_img = cv2.imread(new_user_img_path)

    # ===== 功能6: infer_hairstyle_diy_jy 粗推理 =====
    # 产出 img_res（带新发型的粗推理图）+ 临时文件（user_orig_mask.png, hair_mask_2.png, kpt_1k.txt）
    start = time.time()
    ref_img = cv2.imread(ref_img_path)
    work_dir = os.path.join(userInfo_dir, task_id)
    os.makedirs(work_dir, exist_ok=True)
    import torch
    with torch.no_grad():
        img_res, status, _, landmarks_1k, isEyeOccluded = hairstyle_process.infer_hairstyle_diy_jy(
            origin_img, ref_img, os.path.join(work_dir, task_id), f"{task_id}.png")
    if status != 0:
        raise RuntimeError("发型粗推理失败")
    print(f"[hair_grow] 功能6 粗推理完成: {time.time()-start:.1f}s")

    # ===== 功能8: warpAffine + mask 合并（核心改造点）=====
    start = time.time()
    user_material_dir = os.path.join(work_dir, task_id)
    hair_matting_path = os.path.join(user_material_dir, "hair_mask_2.png")
    user_orig_mask_path = os.path.join(user_material_dir, "user_orig_mask.png")

    origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE)
    new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE)
    if origin_matting is None or new_matting is None:
        raise RuntimeError(f"粗推理产物缺失: origin={origin_matting is None}, new={new_matting is None}")

    # 【关键差异】mask = origin_matting ∪ new_matting ∪ hand_mask
    #   - 不减刘海圆、不减眼眶（与换发型不同）
    #   - 手绘 mask 加入并集，覆盖用户指定的缺发区
    # 手绘 mask 与 origin_matting 同尺寸（都是原图尺寸）
    if hand_mask.shape != origin_matting.shape:
        hand_mask = cv2.resize(hand_mask, (origin_matting.shape[1], origin_matting.shape[0]),
                               interpolation=cv2.INTER_NEAREST)
    _, hand_mask_bin = cv2.threshold(hand_mask, 127, 255, cv2.THRESH_BINARY)
    matting_merge = np.max(
        np.stack([origin_matting, new_matting, hand_mask_bin], axis=2), axis=2
    ).astype(np.uint8)

    # warpAffine 到固定尺寸（与换发型第514-527行完全一致）
    dst_size = (1152, 1536) if is_hr else (576, 768)
    box_info = hairstyle_process.get_body_info(img_res)
    box_w = box_info[2] - box_info[0]
    box_h = box_info[3] - box_info[1]
    scale = min(dst_size[1] / max(box_h, 1), dst_size[0] / max(box_w, 1))
    rotate_center = [(box_info[2] + box_info[0]) * 0.5, (box_info[3] + box_info[1]) * 0.5]
    M = cv2.getRotationMatrix2D(rotate_center, 0, scale)
    M[:, 2] += np.float32([dst_size[0] * 0.5, dst_size[1] * 0.5]) - np.float32(rotate_center)

    crop_result = high_quality_warpAffine(img_res, M, dst_size)
    crop_matting = cv2.warpAffine(matting_merge, M, dst_size)
    mask = (crop_matting > 10).astype(np.float32)
    # dilate（与换发型一致：非HR核3x9，HR核6x18）
    dilate_kernel = (6, 18) if is_hr else (3, 9)
    mask_dilate = cv2.dilate(mask, np.ones(dilate_kernel, np.uint8))
    mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8)
    final_img = crop_result

    # 读取 config.json 的 gender + 训练素材的 prompt tag（与换发型第565-578行一致）
    config_json_path = os.path.join(material_save_path, "config.json")
    with open(config_json_path, "r") as f:
        in_gender = json.load(f)["gender"]

    images_dir = os.path.join(hair_material_dir, "images")
    txt_dir = os.path.join(images_dir, os.listdir(images_dir)[0])
    txt_path = glob.glob(txt_dir + '/*.txt')[0]
    with open(txt_path, 'r') as f:
        p_tag = f.readline()
        if "titor hairstyle, faceless, no human, gray background, simple background" in p_tag:
            p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):]
        else:
            p_tag = ""
    print(f"[hair_grow] 功能8 mask合并+warpAffine完成: {time.time()-start:.1f}s, "
          f"in_gender={in_gender}")

    # ===== 功能9: webui 推理（走 photo_service + LoRA，复用 webui_img2img）=====
    start = time.time()
    sd_result = webui_img2img(
        img=final_img, mask_img=mask_dilate,
        in_gender=in_gender, task_id=task_id,
        hair_id=hair_id, lora_material_path=hair_material_dir,
        tag=p_tag, is_hr=is_hr,
        denoising_strength=0.6, inference_port="57860")
    print(f"[hair_grow] 功能9 webui推理完成: {time.time()-start:.1f}s")

    # ===== 后处理: warpAffine 贴回原图（与换发型第601-605行一致）=====
    M_inv = cv2.invertAffineTransform(M)
    result = origin_img.copy()
    cv2.warpAffine(sd_result, M_inv, (origin_img.shape[1], origin_img.shape[0]),
                   dst=result, borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4)

    # 清理临时文件
    try:
        shutil.rmtree(user_material_dir)
        os.remove(new_user_img_path)
    except Exception:
        pass

    print(f"[hair_grow] 总耗时: {time.time()-start_all:.1f}s")
    return result


def list_hairstyles(hairstyle_dir, train_dir, upload_dir, limit=None):
    """列出所有完整可用的发型。

    参数:
        hairstyle_dir: ref_hairstyle/ 目录
        train_dir: train_material/ 目录
        upload_dir: upload_train_imgs/ 目录
        limit: 限制返回数量（测试用）
    返回:
        [{hair_id, gender, has_lora}, ...]
    """
    styles = []
    if not os.path.isdir(hairstyle_dir):
        return styles
    for hair_id in os.listdir(hairstyle_dir):
        cfg_path = os.path.join(hairstyle_dir, hair_id, "config.json")
        ref_path = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
        lora_path = os.path.join(train_dir, hair_id, "model", "hairstyle_hd_lora.safetensors")
        if not (os.path.exists(cfg_path) and os.path.exists(ref_path) and os.path.exists(lora_path)):
            continue
        try:
            with open(cfg_path) as f:
                gender = json.load(f).get("gender", "unknown")
        except Exception:
            gender = "unknown"
        styles.append({"hair_id": hair_id, "gender": gender})
        if limit and len(styles) >= limit:
            break
    return styles