| """ |
| 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, |
| ) |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| @torch.no_grad() |
| def run_instructpix2pix( |
| pipe, |
| prompt, |
| original_tensor_01, |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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 |
| ) |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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] |
|
|
|
|
| |
| |
| |
|
|
| 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 = [], [] |
| |
| n_first_try_match = 0 |
| n_retry_match = 0 |
| n_no_match = 0 |
|
|
| 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) |
|
|
| |
| 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 |
| |
| 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) |
|
|
| |
| 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() |
