""" 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()