Upload folder using huggingface_hub

blend.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Canny edge mix: edge regions use refine image, "
+            "non-edge regions use weighted average."
+        )
+    )
+    parser.add_argument("--base-dir", type=str, required=True, help="Directory of base images.")
+    parser.add_argument("--refine-dir", type=str, required=True, help="Directory of refine images.")
+    parser.add_argument("--out-dir", type=str, required=True, help="Output root directory.")
+    parser.add_argument(
+        "--exts",
+        type=str,
+        nargs="+",
+        default=["png", "jpg", "jpeg", "bmp", "webp"],
+        help="Image extensions to match (case-insensitive).",
+    )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=0,
+        help="Only process first N matched pairs (0 means all).",
+    )
+
+    # Keep these five tunable parameters only.
+    parser.add_argument(
+        "--nonedge-alpha",
+        type=float,
+        default=0.5,
+        help="Base-image weight for non-edge averaging (0~1).",
+    )
+    parser.add_argument("--canny-low", type=int, default=80, help="Canny lower threshold.")
+    parser.add_argument("--canny-high", type=int, default=160, help="Canny upper threshold.")
+    parser.add_argument(
+        "--canny-dilate",
+        type=int,
+        default=1,
+        help="Dilate iterations for edge mask.",
+    )
+    parser.add_argument(
+        "--edge-feather",
+        type=int,
+        default=1,
+        help="Gaussian feather radius for edge mask (0 means hard edge).",
+    )
+    return parser.parse_args()
+
+
+def build_pairs(base_dir: Path, refine_dir: Path, exts):
+    exts_set = {e.lower().lstrip(".") for e in exts}
+    base_map = {}
+    for p in base_dir.rglob("*"):
+        if p.is_file() and p.suffix.lower().lstrip(".") in exts_set:
+            rel = p.relative_to(base_dir).as_posix()
+            base_map[rel] = p
+
+    pairs = []
+    missing = 0
+    for rel, base_path in sorted(base_map.items()):
+        refine_path = refine_dir / rel
+        if refine_path.exists():
+            pairs.append((rel, base_path, refine_path))
+        else:
+            missing += 1
+    return pairs, missing
+
+
+def canny_edge_mask(base_bgr_u8, canny_low, canny_high, dilate_iter=1, feather_radius=1):
+    gray = cv2.cvtColor(base_bgr_u8, cv2.COLOR_BGR2GRAY)
+    gray = cv2.GaussianBlur(gray, (5, 5), 0.0)
+    edge = cv2.Canny(gray, int(canny_low), int(canny_high))
+    if dilate_iter > 0:
+        kernel = np.ones((3, 3), np.uint8)
+        edge = cv2.dilate(edge, kernel, iterations=int(dilate_iter))
+    edge01 = edge.astype(np.float32) / 255.0
+    if feather_radius > 0:
+        k = int(feather_radius) * 2 + 1
+        edge01 = cv2.GaussianBlur(edge01, (k, k), 0.0)
+    edge01 = np.clip(edge01, 0.0, 1.0)
+    return edge01, edge
+
+
+def blend_with_base_weight(base_bgr_u8, refine_bgr_u8, base_weight_01):
+    base = base_bgr_u8.astype(np.float32) / 255.0
+    refine = refine_bgr_u8.astype(np.float32) / 255.0
+    w = base_weight_01[:, :, None]
+    out = w * base + (1.0 - w) * refine
+    out = np.clip(out * 255.0 + 0.5, 0.0, 255.0).astype(np.uint8)
+    return out
+
+
+def ensure_parent(path: Path):
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+
+def main():
+    args = parse_args()
+    base_dir = Path(args.base_dir)
+    refine_dir = Path(args.refine_dir)
+    out_dir = Path(args.out_dir)
+
+    out_blend = out_dir / "blended"
+    out_mask = out_dir / "mask01"
+    out_mask_vis = out_dir / "mask_vis"
+    out_blend.mkdir(parents=True, exist_ok=True)
+    out_mask.mkdir(parents=True, exist_ok=True)
+    out_mask_vis.mkdir(parents=True, exist_ok=True)
+
+    pairs, missing = build_pairs(base_dir, refine_dir, args.exts)
+    if args.limit > 0:
+        pairs = pairs[: args.limit]
+    print(f"Found pairs: {len(pairs)}, missing in refine: {missing}")
+    if not pairs:
+        raise RuntimeError("No matching image pairs found.")
+
+    fail_count = 0
+    alpha_nonedge = float(np.clip(args.nonedge_alpha, 0.0, 1.0))
+    for idx, (rel, base_path, refine_path) in enumerate(pairs, 1):
+        base = cv2.imread(str(base_path), cv2.IMREAD_COLOR)
+        refine = cv2.imread(str(refine_path), cv2.IMREAD_COLOR)
+        if base is None or refine is None:
+            print(f"[WARN] Failed to read: {rel}")
+            fail_count += 1
+            continue
+        if base.shape != refine.shape:
+            print(f"[WARN] Shape mismatch, skip: {rel}, {base.shape} vs {refine.shape}")
+            fail_count += 1
+            continue
+
+        edge01, edge_bin = canny_edge_mask(
+            base,
+            args.canny_low,
+            args.canny_high,
+            dilate_iter=args.canny_dilate,
+            feather_radius=args.edge_feather,
+        )
+
+        # Edge: use refine directly (base weight = 0).
+        # Non-edge: weighted average by nonedge-alpha.
+        base_weight = np.clip((1.0 - edge01) * alpha_nonedge, 0.0, 1.0)
+        out = blend_with_base_weight(base, refine, base_weight)
+
+        rel_path = Path(rel)
+        p_blend = out_blend / rel_path
+        p_mask = out_mask / rel_path
+        p_mvis = out_mask_vis / rel_path
+        ensure_parent(p_blend)
+        ensure_parent(p_mask)
+        ensure_parent(p_mvis)
+
+        cv2.imwrite(str(p_blend), out)
+        cv2.imwrite(str(p_mask), (base_weight * 255.0 + 0.5).astype(np.uint8))
+
+        heat = cv2.applyColorMap(edge_bin.astype(np.uint8), cv2.COLORMAP_JET)
+        vis = cv2.addWeighted(refine, 0.78, heat, 0.22, 0.0)
+        cv2.imwrite(str(p_mvis), vis)
+
+        if idx % 100 == 0 or idx == len(pairs):
+            print(f"Processed {idx}/{len(pairs)}")
+
+    print(f"Done. failures={fail_count}, output={out_dir}")
+
+
+if __name__ == "__main__":
+    main()

infer_ir.py

@@ -17,7 +17,7 @@ weight_dtype = torch.bfloat16  # 可改为 torch.float16 / torch.float32
 
 vae_path = BASE_DIR / "vae.pth"
 transformer_path = BASE_DIR / "transformer.pth"
-unet_path = BASE_DIR / "unet.pth"
+unet_path = BASE_DIR / "unet_178.pth"
 text_embedding_path = BASE_DIR / "text_embeddings.pth"
 
 output_folder = BASE_DIR / "output"
@@ -111,6 +111,13 @@ if not input_folder.exists() or not input_folder.is_dir():
     raise ValueError(f"Invalid input folder: {input_folder}")
 
 output_folder.mkdir(parents=True, exist_ok=True)
+
+# 新增两个子文件夹
+x_pred_folder = output_folder / "x_pred"
+x_pred_final_folder = output_folder / "x_pred_final"
+x_pred_folder.mkdir(parents=True, exist_ok=True)
+x_pred_final_folder.mkdir(parents=True, exist_ok=True)
+
 image_paths = collect_image_paths(input_folder)
 
 shift_factor = getattr(vae.config, "shift_factor", 0.0)
@@ -167,9 +174,16 @@ for img_path in tqdm(image_paths, desc="Processing"):
     x_pred_final = x_pred + unet_pred
 
     # ===== 保存 =====
-    out_pil = vae_output_to_pil(x_pred_final)
     out_name = img_path.stem + ".png"
-    out_path = output_folder / out_name
-    out_pil.save(out_path, format="PNG", compress_level=0)
+
+    # 保存 x_pred
+    x_pred_pil = vae_output_to_pil(x_pred)
+    x_pred_path = x_pred_folder / out_name
+    x_pred_pil.save(x_pred_path, format="PNG", compress_level=0)
+
+    # 保存 x_pred_final
+    x_pred_final_pil = vae_output_to_pil(x_pred_final)
+    x_pred_final_path = x_pred_final_folder / out_name
+    x_pred_final_pil.save(x_pred_final_path, format="PNG", compress_level=0)
 
 print(f"Done. Saved to: {output_folder}")

read_from_dir.py

@@ -74,7 +74,7 @@ def process_images_from_dir(input_dir, output_base_dir=None):
 
 if __name__ == "__main__":
     # 配置变量
-    input_directory = "./output"  # 修改为你的输入目录路径
+    input_directory = "./final_result/blended"  # 修改为你的输入目录路径
     output_base_directory = "./"  # 如果为None，输出目录会在输入目录的父目录下创建
     
     process_images_from_dir(input_directory, output_base_directory)

readme.md

@@ -8,3 +8,6 @@ pip install -r requirements.txt -i https://mirrors.cloud.tencent.com/pypi/simple
 
 ### 按照run.sh命令运行即可
 cd到当前项目路径，运行run.sh即可
+最后的结果在processed_jpg目录下
+
+如果运行出现问题，可以联系：chenyx.cs@gmail.com

requirements.txt

@@ -1,3 +1,4 @@
 diffusers==0.34.0
 torch==2.4.1
-torchvision==0.19.1
+torchvision==0.19.1
+opencv-python

run.sh

@@ -1,7 +1,12 @@
 # input_folder 换成本地的输入路径,结果产生在./output
 python ./infer_ir.py \
-  --input_folder /mnt/tidal-sh01/dataset/media_image_algo_tidalfs/zhongyilian/Datasets/Test/LoViF_LQ
+  --input_folder xxxxxxxxxxxxxxx
 
-# 将png结果转为比赛要求的jpg
-python read_from_dir.py
+python ./blend.py \
+  --base-dir ./output/x_pred \
+  --refine-dir ./output/x_pred_final \
+  --out-dir ./final_result \
+  --nonedge-alpha 0.5 --canny-low 80 --canny-high 160 --canny-dilate 1 --edge-feather 1
 
+# 将png结果转为比赛要求的jpg
+python read_from_dir.py

unet_178.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a70f6ea32f577af5573612956d69b3101444fae98575ba6835e10a56410b2e2
+size 454870132