Blame for text2layout/generate_layout_rib.py

"""
Inference for InstructPix2Pix trained with MyDataset_rib.

关键点:
- 训练时 input(original) = stack([empty, organ_norm, rib_norm]), 值域 [0,1]
- 训练时 target(edited)  = stack([disease_bin, organ_norm, rib_norm]), 值域 [0,1]
- VAE 期望输入归一化到 [-1, 1] (pipeline 内部会做, 但我们绕开默认 image_processor)
- 推理时从生成图的第 0 通道提取 disease mask, 然后二值化 + 形态学清理
- 注意 organ / rib 的 label 标签 (60/120/180/240 等) 在训练时被 _remap_to_sequential 重映射
  并归一化到 [0,1], inference 必须用完全相同的方式构造输入
- 每次 retry 都会打印 [MATCH]/[DIFF] 对比, 看模型实际生成的位置/严重度
"""

import os
import re
import json
import copy
import argparse
import torch
import numpy as np
from PIL import Image
from tqdm import tqdm
from skimage.morphology import disk
from scipy.ndimage import binary_erosion, binary_dilation

from diffusers import (
    StableDiffusionInstructPix2PixPipeline,
    UNet2DConditionModel,
    AutoencoderKL,
    DDIMScheduler,
)


# ============================================================
# 复刻 MyDataset_rib 中的工具函数 (保持完全一致!)
# ============================================================

def _load_gray_int(path):
    arr = np.array(Image.open(path).convert("L"))
    return arr.astype(np.int32)


def _remap_to_sequential(label_map_int):
    vals = np.unique(label_map_int)
    vals = vals[vals != 0]
    if len(vals) == 0:
        return label_map_int.astype(np.int32), 0
    out = np.zeros_like(label_map_int, dtype=np.int32)
    for i, v in enumerate(vals, start=1):
        out[label_map_int == v] = i
    return out, int(out.max())


def _norm(ch, maxv):
    return ch.astype(np.float32) / float(maxv) if maxv > 0 else ch.astype(np.float32)


def _resize_label_map(label_map_int, target_hw):
    """label map 用 NEAREST resize, 保持离散标签"""
    H, W = target_hw
    im = Image.fromarray(label_map_int.astype(np.int32), mode="I")
    if im.size != (W, H):
        im = im.resize((W, H), resample=Image.NEAREST)
    return np.array(im).astype(np.int32)


def build_original_tensor(organ_path, rib_path, resolution=512):
    """
    严格复刻 MyDataset_rib.__getitem__ 中 original 的构造逻辑
    返回: torch.Tensor (3, H, W), 值域 [0, 1]
    """
    organ_map_raw = _load_gray_int(organ_path)
    organ_map_raw = _resize_label_map(organ_map_raw, (resolution, resolution))
    organ_map, organ_max = _remap_to_sequential(organ_map_raw)

    rib_map_raw = _load_gray_int(rib_path)
    rib_map_raw = _resize_label_map(rib_map_raw, (resolution, resolution))
    rib_map, rib_max = _remap_to_sequential(rib_map_raw)

    H, W = organ_map.shape
    empty_ch = np.zeros((H, W), dtype=np.float32)
    organ_ch = _norm(organ_map, organ_max)
    rib_ch = _norm(rib_map, rib_max)

    original = np.stack([empty_ch, organ_ch, rib_ch], axis=0).astype(np.float32)
    return torch.from_numpy(original)  # (3, H, W) in [0, 1]


# ============================================================
# Inference 主流程: 绕开默认 image_processor, 直接喂 VAE
# ============================================================

@torch.no_grad()
def run_instructpix2pix(
    pipe,
    prompt,
    original_tensor_01,           # (3, H, W) in [0, 1]
    num_inference_steps=50,
    guidance_scale=7.0,
    image_guidance_scale=1.5,
    generator=None,
    device="cuda",
    dtype=torch.float16,
):
    """
    手动实现 InstructPix2Pix 的 inference loop, 完全控制 image 编码方式,
    保证和训练分布一致 (不经过 pipeline 默认的 [-1,1] 归一化两次).
    """
    vae = pipe.vae
    unet = pipe.unet
    text_encoder = pipe.text_encoder
    tokenizer = pipe.tokenizer
    scheduler = pipe.scheduler  # 用 pipeline 自带 scheduler

    # ---- 1. 编码 prompt ----
    text_inputs = tokenizer(
        [prompt],
        padding="max_length",
        max_length=tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    ).input_ids.to(device)
    prompt_emb = text_encoder(text_inputs)[0]

    uncond_inputs = tokenizer(
        [""],
        padding="max_length",
        max_length=tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    ).input_ids.to(device)
    uncond_emb = text_encoder(uncond_inputs)[0]

    text_emb_all = torch.cat([prompt_emb, uncond_emb, prompt_emb], dim=0).to(dtype)

    # ---- 2. 编码 image ----
    image = original_tensor_01.unsqueeze(0).to(device=device, dtype=dtype)
    image_latents = vae.encode(image).latent_dist.mode()
    uncond_image_latents = torch.zeros_like(image_latents)
    image_latents_all = torch.cat(
        [image_latents, image_latents, uncond_image_latents], dim=0
    )

    # ---- 3. 初始化 latents ----
    scheduler.set_timesteps(num_inference_steps, device=device)
    timesteps = scheduler.timesteps
    latent_shape = (1, 4, image.shape[-2] // 8, image.shape[-1] // 8)
    latents = torch.randn(latent_shape, generator=generator, device=device, dtype=dtype)
    latents = latents * scheduler.init_noise_sigma

    # ---- 4. Denoise loop ----
    for t in timesteps:
        latent_model_input = torch.cat([latents] * 3, dim=0)
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
        latent_model_input = torch.cat([latent_model_input, image_latents_all], dim=1)

        noise_pred = unet(
            latent_model_input, t, encoder_hidden_states=text_emb_all
        ).sample
        noise_pred_cond, noise_pred_text_uncond, noise_pred_image_uncond = noise_pred.chunk(3)

        noise_pred = (
            noise_pred_image_uncond
            + image_guidance_scale * (noise_pred_text_uncond - noise_pred_image_uncond)
            + guidance_scale * (noise_pred_cond - noise_pred_text_uncond)
        )
        latents = scheduler.step(noise_pred, t, latents).prev_sample

    # ---- 5. 解码 ----
    latents_for_decode = latents / vae.config.scaling_factor
    image_out = vae.decode(latents_for_decode).sample
    image_out = image_out.clamp(0, 1)
    return image_out[0]


# ============================================================
# Mask 提取 + prompt 一致性校验
# ============================================================

def extract_mask_from_first_channel(image_chw_tensor, threshold=0.5, morph_radius=3, morph_mode="open"):
    first_channel = image_chw_tensor[0].detach().cpu().float().numpy()
    binary_mask = first_channel > threshold

    selem = disk(morph_radius)
    if morph_mode == "open":
        binary_mask = binary_erosion(binary_mask, structure=selem)
        binary_mask = binary_dilation(binary_mask, structure=selem)
    else:
        binary_mask = binary_dilation(binary_mask, structure=selem)
        binary_mask = binary_erosion(binary_mask, structure=selem)

    binary_mask = (binary_mask.astype(np.uint8)) * 255
    return Image.fromarray(binary_mask, mode="L")


def pil_to_gray_np(img):
    if isinstance(img, Image.Image):
        return np.array(img.convert("L"))
    arr = np.array(img)
    if arr.ndim == 3:
        arr = arr[..., 0]
    return arr


def define_organ_parts(mask):
    left_lung_coords = np.where(mask == 60)
    right_lung_coords = np.where(mask == 120)

    if left_lung_coords[0].size > 0:
        left_min, left_max = left_lung_coords[0].min(), left_lung_coords[0].max()
        left_upper_boundary = left_min + (left_max - left_min) // 3
        left_middle_boundary = left_min + 2 * (left_max - left_min) // 3
    else:
        left_upper_boundary, left_middle_boundary = 0, 0

    if right_lung_coords[0].size > 0:
        right_min, right_max = right_lung_coords[0].min(), right_lung_coords[0].max()
        right_upper_boundary = right_min + (right_max - right_min) // 3
        right_middle_boundary = right_min + 2 * (right_max - right_min) // 3
    else:
        right_upper_boundary, right_middle_boundary = 0, 0

    rows = np.arange(mask.shape[0])[:, None]
    organ_parts = {
        "left upper lung": (mask == 60) & (rows <= left_upper_boundary),
        "left middle lung": (mask == 60) & ((rows > left_upper_boundary) & (rows <= left_middle_boundary)),
        "left lower lung": (mask == 60) & (rows > left_middle_boundary),
        "right upper lung": (mask == 120) & (rows <= right_upper_boundary),
        "right middle lung": (mask == 120) & ((rows > right_upper_boundary) & (rows <= right_middle_boundary)),
        "right lower lung": (mask == 120) & (rows > right_middle_boundary),
        "heart": (mask == 180),
        "mediastinum": (mask == 240),
    }
    return organ_parts


def calculate_width(region_mask):
    non_zero_columns = np.where(region_mask > 0)[1]
    if len(non_zero_columns) == 0:
        return 0
    return non_zero_columns.max() - non_zero_columns.min() + 1


def process_organ_and_mask_from_array(disease, organ, mask):
    organ_parts = define_organ_parts(organ)
    overlap_results = {}
    for part, mask_part in organ_parts.items():
        overlap_area = np.sum((mask_part > 0) & (mask > 0))
        if overlap_area > 0:
            overlap_results[part] = overlap_area
    if not overlap_results:
        return None
    main_part = max(overlap_results, key=overlap_results.get)

    if disease == "Cardiomegaly":
        if main_part == "heart":
            organ_width = calculate_width(organ)
            mask_width = calculate_width(mask)
            if organ_width == 0:
                return None
            ratio = mask_width / organ_width
            if ratio <= 0.55:
                severity = "mild"
            elif ratio < 0.6:
                severity = "moderate"
            else:
                severity = "severe"
            return disease, main_part, severity
        return None

    if disease == "Enlarged Cardiomediastinum":
        organ_width = calculate_width(organ)
        mask_width = calculate_width(mask)
        if organ_width == 0 or mask_width == 0:
            return None
        ratio = mask_width / organ_width
        if ratio <= 0.55:
            severity = "mild"
        elif 0.55 < ratio < 0.6:
            severity = "moderate"
        else:
            severity = "severe"
        return disease, "heart and mediastinum", severity

    if overlap_results[main_part] > np.sum(mask > 0) * 0.5:
        location_label = main_part
        severity = "mild"
    else:
        left_regions = {"left upper lung", "left middle lung", "left lower lung"}
        right_regions = {"right upper lung", "right middle lung", "right lower lung"}
        active_regions = set(overlap_results.keys())

        left_lung = (organ == 60)
        left_lung_area = np.sum(left_lung)
        left_overlap_area = np.sum((mask > 0) & left_lung)
        left_lung_ratio = left_overlap_area / left_lung_area if left_lung_area > 0 else 0.0

        right_lung = (organ == 120)
        right_lung_area = np.sum(right_lung)
        right_overlap_area = np.sum((mask > 0) & right_lung)
        right_lung_ratio = right_overlap_area / right_lung_area if right_lung_area > 0 else 0.0

        left_overlap = active_regions & left_regions
        right_overlap = active_regions & right_regions
        if left_overlap and right_overlap:
            location_label = "bilateral lung"
        elif left_overlap:
            location_label = "left lung"
        elif right_overlap:
            location_label = "right lung"
        else:
            location_label = None

        if disease == "Support Devices":
            severity = None
        else:
            if left_lung_ratio < 0.3 and right_lung_ratio < 0.3:
                severity = "mild"
            elif (left_lung_ratio > 0.6 or right_lung_ratio > 0.6 or
                  (left_lung_ratio + right_lung_ratio) > 0.6):
                severity = "severe"
            else:
                severity = "moderate"
    return disease, location_label, severity




def build_prompt_from_mask(disease, organ_path, mask_pil):
    organ_raw = _load_gray_int(organ_path)
    H, W = organ_raw.shape

    mask_pil_resized = mask_pil.resize((W, H), resample=Image.NEAREST)
    mask_np = pil_to_gray_np(mask_pil_resized)
    mask_np = (mask_np > 0).astype(np.uint8) * 255

    result = process_organ_and_mask_from_array(disease, organ_raw, mask_np)
    if result is None:
        return None
    disease_calc, loc, severity = result
    if loc is not None and severity is not None:
        return f"A Chest X-ray semantic mask with {severity} {disease_calc} on {loc}."
    elif loc is not None:
        return f"A Chest X-ray semantic mask with {disease_calc} on {loc}."
    else:
        return disease_calc


def normalize_prompt_text(text):
    if text is None:
        return ""
    text = re.sub(r"\s+", " ", text.strip())
    text = re.sub(r"\s*,\s*", ", ", text)
    return text.lower()


def prompts_match(a, b):
    return normalize_prompt_text(a) == normalize_prompt_text(b)


def make_generated_mask_relpath(file_name, disease):
    base_dir = os.path.dirname(file_name)
    stem = os.path.splitext(os.path.basename(file_name))[0]
    safe_disease = disease.replace("/", "_")
    return os.path.join(base_dir, f"{stem}_{safe_disease}.png")


def replace_attn_list_with_generated_mask(attn_list, generated_mask_relpath):
    return [[d[0], generated_mask_relpath] for d in attn_list]


# ============================================================
# 主入口
# ============================================================

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--output_dir", type=str, required=True)
    p.add_argument("--model_dir", type=str, required=True,
                   help="包含 checkpoint-XXXX 的训练输出目录")
    p.add_argument("--checkpoint_num", type=int, required=True)
    p.add_argument("--pretrained_model_name_or_path", type=str,
                   default="CompVis/stable-diffusion-v1-4")
    p.add_argument("--image_root", type=str, required=True,
                   help="organ/rib 等相对路径的 base dir (即训练时的 train_data_dir)")
    p.add_argument("--test_json_path", type=str, required=True)
    p.add_argument("--resolution", type=int, default=512)
    p.add_argument("--max_attempts", type=int, default=5)
    p.add_argument("--num_inference_steps", type=int, default=50)
    p.add_argument("--guidance_scale", type=float, default=7.0)
    p.add_argument("--image_guidance_scale", type=float, default=1.5)
    p.add_argument("--mask_threshold", type=float, default=0.5,
                   help="第 0 通道二值化阈值 (输出已是 [0,1], 用 0.5)")
    p.add_argument("--morph_radius", type=int, default=3)
    p.add_argument("--morph_mode", type=str, default="open", choices=["open", "close"])
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--save_full_rgb", action="store_true",
                   help="同时保存生成的 3 通道图 (debug 用)")
    return p.parse_args()


def load_pipeline(args, device, dtype):
    ckpt_path = os.path.join(args.model_dir, f"checkpoint-{args.checkpoint_num}")
    print(f"Loading base pipeline from {args.pretrained_model_name_or_path}")
    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        args.pretrained_model_name_or_path,
        torch_dtype=dtype,
        safety_checker=None,
    )
    print(f"Loading fine-tuned UNet from {ckpt_path}")
    pipe.unet = UNet2DConditionModel.from_pretrained(
        ckpt_path, subfolder="unet", torch_dtype=dtype
    )
    pipe.set_progress_bar_config(disable=True)
    pipe.to(device)
    return pipe


def load_items(json_path):
    items = []
    with open(json_path, "r") as f:
        for line in f:
            line = line.strip()
            if line:
                items.append(json.loads(line))
    return items


def _resolve(rel_or_abs, base_dir):
    return rel_or_abs if os.path.isabs(rel_or_abs) else os.path.join(base_dir, rel_or_abs.lstrip("/"))


def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.float16 if device.type == "cuda" else torch.float32

    pipe = load_pipeline(args, device, dtype)
    os.makedirs(args.output_dir, exist_ok=True)

    items = load_items(args.test_json_path)
    print(f"Total images: {len(items)}")

    generator = None
    if args.seed is not None:
        generator = torch.Generator(device=device).manual_seed(args.seed)

    success_records, failed_records = [], []
    # 统计 retry 的分布
    n_first_try_match = 0   # 第 1 次就 match
    n_retry_match = 0       # retry 之后 match
    n_no_match = 0          # max_attempts 都没 match

    for idx, item in enumerate(tqdm(items, desc="Generating masks")):
        raw_item = copy.deepcopy(item)
        file_name = raw_item["file_name"]
        original_prompt = raw_item["prompt"]
        attn_list = raw_item.get("attn_list", [])
        organ_rel = raw_item["organ"]
        rib_rel = raw_item.get("rib", "")
        organ_path = _resolve(organ_rel, args.image_root)
        rib_path = _resolve(rib_rel, args.image_root)

        if len(attn_list) == 0:
            failed_item = copy.deepcopy(raw_item)
            failed_item["generated_prompt"] = None
            failed_item["reason"] = "empty attn_list"
            failed_records.append(failed_item)
            continue

        disease = attn_list[0][0]
        generated_mask_relpath = make_generated_mask_relpath(file_name, disease)
        generated_mask_save_path = os.path.join(args.output_dir, generated_mask_relpath)
        os.makedirs(os.path.dirname(generated_mask_save_path), exist_ok=True)

        # Normal: 直接空 mask
        if disease == "Normal":
            organ_pil = Image.open(organ_path).convert("L")
            empty_mask = Image.new("L", organ_pil.size, 0)
            empty_mask.save(generated_mask_save_path)
            success_item = copy.deepcopy(raw_item)
            success_item["attn_list"] = replace_attn_list_with_generated_mask(
                attn_list, generated_mask_relpath
            )
            success_records.append(success_item)
            tqdm.write(
                f"[{idx+1:>5}/{len(items)}] [NORMAL] {os.path.basename(file_name)}  "
                f"(empty mask, skip generation)"
            )
            continue

        original_tensor = build_original_tensor(
            organ_path, rib_path, resolution=args.resolution
        )

        success = False
        last_generated_prompt = None
        # 头部打印一次 original
        tqdm.write(
            f"[{idx+1:>5}/{len(items)}] {os.path.basename(file_name)}  disease={disease}\n"
            f"    original : {original_prompt}"
        )

        for attempt in range(args.max_attempts):
            generated_chw = run_instructpix2pix(
                pipe,
                prompt=original_prompt,
                original_tensor_01=original_tensor,
                num_inference_steps=args.num_inference_steps,
                guidance_scale=args.guidance_scale,
                image_guidance_scale=args.image_guidance_scale,
                generator=generator,
                device=device,
                dtype=dtype,
            )

            satisfied_mask = extract_mask_from_first_channel(
                generated_chw,
                threshold=args.mask_threshold,
                morph_radius=args.morph_radius,
                morph_mode=args.morph_mode,
            )

            generated_prompt = build_prompt_from_mask(
                disease=disease,
                organ_path=organ_path,
                mask_pil=satisfied_mask,
            )
            last_generated_prompt = generated_prompt

            is_match = (generated_prompt is not None
                        and prompts_match(original_prompt, generated_prompt))
            tag = "MATCH" if is_match else ("NONE " if generated_prompt is None else "DIFF ")
            tqdm.write(
                f"    try {attempt+1}/{args.max_attempts} [{tag}]: {generated_prompt}"
            )

            if generated_prompt is None:
                continue
            if not is_match:
                continue

            # 成功
            satisfied_mask.save(generated_mask_save_path)

            if args.save_full_rgb:
                rgb_np = (generated_chw.detach().cpu().float().numpy()
                          .transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
                Image.fromarray(rgb_np, mode="RGB").save(
                    generated_mask_save_path.replace(".png", "_rgb.png")
                )

            success_item = copy.deepcopy(raw_item)
            success_item["attn_list"] = replace_attn_list_with_generated_mask(
                attn_list, generated_mask_relpath
            )
            success_item["generated_prompt"] = generated_prompt
            success_item["attempts_used"] = attempt + 1
            success_records.append(success_item)

            if attempt == 0:
                n_first_try_match += 1
            else:
                n_retry_match += 1
            success = True
            break

        if not success:
            n_no_match += 1
            tqdm.write(
                f"    [FAIL] after {args.max_attempts} attempts. "
                f"last={last_generated_prompt}"
            )
            failed_item = copy.deepcopy(raw_item)
            failed_item["generated_prompt"] = last_generated_prompt
            failed_item["reason"] = "prompt mismatch or invalid mask"
            failed_records.append(failed_item)

    # 保存 jsonl
    base = os.path.basename(args.test_json_path)
    success_json_path = os.path.join(args.output_dir, base)
    failed_json_path = os.path.join(args.output_dir, f"failed_{base}")

    with open(success_json_path, "w") as f:
        for r in success_records:
            f.write(json.dumps(r) + "\n")
    with open(failed_json_path, "w") as f:
        for r in failed_records:
            f.write(json.dumps(r) + "\n")

    n = len(items)
    fr = len(failed_records) / n * 100 if n else 0.0
    n_attempt_total = n_first_try_match + n_retry_match + n_no_match
    print(f"\n===== Summary =====")
    print(f"Total: {n}, Success: {len(success_records)}, Failed: {len(failed_records)}, "
          f"Failure Rate: {fr:.2f}%")
    if n_attempt_total > 0:
        print(f"Among {n_attempt_total} non-Normal samples:")
        print(f"  first-try match : {n_first_try_match}  ({n_first_try_match/n_attempt_total*100:.1f}%)")
        print(f"  retry  match    : {n_retry_match}  ({n_retry_match/n_attempt_total*100:.1f}%)")
        print(f"  no match (fail) : {n_no_match}  ({n_no_match/n_attempt_total*100:.1f}%)")
    print(f"Success jsonl: {success_json_path}")
    print(f"Failed  jsonl: {failed_json_path}")


if __name__ == "__main__":
    main()