Upload 17 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo/video_3.mp4 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,209 @@
+import os
+import cv2
+import torch
+import face_alignment
+import numpy as np
+import mediapipe as mp
+import tempfile
+from multiprocessing import cpu_count
+from PIL import Image
+import streamlit as st
+from test import test_single_image
+
+
+reference_heatmap_dir = "motion_transfer/dataset_single/reference_heatmap"
+output_dir = r"motion_transfer\dataset_single\test_heatmap"
+final_output = "motion_transfer/outputs/final_result.mp4"
+os.makedirs(output_dir, exist_ok=True)
+
+num_workers = min(cpu_count(), 4)
+target_size = 256
+SIGMA = 2.0
+NUM_FACE_POINTS = 68
+
+
+fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D,
+                                  device="cuda" if torch.cuda.is_available() else "cpu")
+mp_face_detection = mp.solutions.face_detection
+detector = mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
+
+
+
+def gaussian_heatmaps(points, H, W, sigma=2.0):
+    yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
+    heatmaps = np.exp(-((xx[..., None] - points[:, 0]) ** 2 +
+                        (yy[..., None] - points[:, 1]) ** 2) / (2 * sigma ** 2))
+    return heatmaps.astype(np.float32)
+
+
+def extract_keypoints(hmap):
+    kps = []
+    for i in range(hmap.shape[2]):
+        y, x = np.where(hmap[:, :, i] > 0)
+        if len(x) > 0:
+            kps.append([np.mean(x), np.mean(y)])
+        else:
+            kps.append([0, 0])
+    return np.array(kps, dtype=np.float32)
+
+
+def trim_video(input_path, output_path, max_seconds=7):
+    cap = cv2.VideoCapture(input_path)
+    if not cap.isOpened():
+        print("Error opening video")
+        return False
+    fps = int(round(cap.get(cv2.CAP_PROP_FPS)))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    max_frames = min(total_frames, fps * max_seconds)
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    count = 0
+    while cap.isOpened() and count < max_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        out.write(frame)
+        count += 1
+    cap.release()
+    out.release()
+    return True
+
+
+def crop_head_with_bg(img_rgb, target_size=256, margin_top=0.6, margin_sides=0.3, margin_bottom=0.4):
+
+    ih, iw, _ = img_rgb.shape
+    results = detector.process(img_rgb)
+    if not results.detections:
+        return None
+
+    det = results.detections[0]
+    bbox = det.location_data.relative_bounding_box
+    x1 = int(bbox.xmin * iw)
+    y1 = int(bbox.ymin * ih)
+    w = int(bbox.width * iw)
+    h = int(bbox.height * ih)
+
+   
+    x1 = max(0, int(x1 - w * margin_sides))
+    x2 = min(iw, int(x1 + w * (1 + 2 * margin_sides)))
+    y1 = max(0, int(y1 - h * margin_top))
+    y2 = min(ih, int(y1 + h * (1 + margin_bottom + margin_top)))
+
+    cropped = img_rgb[y1:y2, x1:x2]
+    ch, cw = cropped.shape[:2]
+
+
+    scale = target_size / max(ch, cw)
+    new_w, new_h = int(cw * scale), int(ch * scale)
+    resized = cv2.resize(cropped, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
+
+
+    blurred_bg = cv2.GaussianBlur(resized, (51, 51), 0)
+    background = cv2.resize(blurred_bg, (target_size, target_size), interpolation=cv2.INTER_AREA)
+
+    y_offset = (target_size - new_h) // 2
+    x_offset = (target_size - new_w) // 2
+    background[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized
+
+    return background
+
+
+# Streamlit
+st.title("Sketch to Live")
+
+src_img = st.file_uploader("Upload face sketch", type=["jpg", "png"])
+cropped_head = None
+
+if src_img is not None:
+    
+    pil_img = Image.open(src_img).convert("RGB")
+    img_rgb = np.array(pil_img)
+    ih, iw, _ = img_rgb.shape
+    st.write(f"Uploaded image size: {iw}×{ih}")
+
+    if ih < target_size or iw < target_size:
+        st.warning("Image too small ({iw}×{ih}).Please upload one larger than {target_size}×{target_size}.")
+    else:
+        cropped_head = crop_head_with_bg(img_rgb, target_size=target_size)
+        if cropped_head is None:
+            st.warning("No face detected. Try another image.")
+        else:
+            st.subheader("Face Preview")
+            st.image(
+                cropped_head,
+                caption="Cropped Head",
+                width=256,
+                channels="RGB",
+                output_format="PNG",
+            )
+
+            # Save 
+            cv2.imwrite("cropped_head.png",
+                        cv2.cvtColor(cropped_head, cv2.COLOR_RGB2BGR),
+                        [cv2.IMWRITE_PNG_COMPRESSION, 0])
+
+
+
+if st.button("Lively Sketch"):
+    if cropped_head is None:
+        st.error("Please upload a face image.")
+    else:
+        progress_text = st.empty()
+        progress_bar = st.progress(0)
+        frame_preview = st.empty()
+        progress_text.text("Processing")
+
+        H, W = cropped_head.shape[:2]
+        fa_out = fa.get_landmarks(cropped_head)
+        if fa_out is None or len(fa_out) == 0:
+            st.error(" No face landmarks detected.")
+        else:
+            face68 = fa_out[0].astype(np.float32)
+            single_heatmap = gaussian_heatmaps(face68, H, W, sigma=SIGMA)
+            single_kp = face68
+
+            ref_files = sorted([f for f in os.listdir(reference_heatmap_dir) if f.endswith(".npy")])
+            if len(ref_files) == 0:
+                st.error(" No reference heatmaps found!")
+            else:
+                ref_heatmaps = [np.load(os.path.join(reference_heatmap_dir, f)) for f in ref_files]
+                ref_kp_list = [extract_keypoints(hm) for hm in ref_heatmaps]
+                ref_base_kp = ref_kp_list[0]
+                motion_vectors = [kp - ref_base_kp for kp in ref_kp_list]
+
+                os.makedirs(output_dir, exist_ok=True)
+                total_frames = len(motion_vectors)
+
+                for frame_idx, displacement in enumerate(motion_vectors):
+                    moved_kp = single_kp + displacement
+                    new_heatmap = gaussian_heatmaps(moved_kp, H, W, sigma=SIGMA)
+                    np.save(os.path.join(output_dir, f"{frame_idx:05d}.npy"), new_heatmap)
+
+                    frame_preview.image(cropped_head, width=128)
+                    progress_bar.progress(int((frame_idx + 1) / total_frames * 100))
+
+                temp_img_path = "cropped_head.png"
+                test_single_image(temp_img_path, output_dir, final_output)
+
+                trimmed_output = "trimmed_result.mp4"
+                trim_video(final_output, trimmed_output, max_seconds=7)
+
+                progress_bar.progress(100)
+                progress_text.text("Done!")
+                frame_preview.empty()
+
+                st.success("Sketch-to-Live")
+                with open(trimmed_output, "rb") as f:
+                    st.download_button("Download Result Video", f, file_name="Sketch.mp4")
+
+
+
+st.markdown("""
+<div style='position: fixed; bottom: 10px; right: 10px; color: gray; font-size: 12px;'>
+    <b>Inspired by prior explicit motion transfer methods.</b>
+</div>
+""", unsafe_allow_html=True)

checkpoints/best.pth

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

checkpoints/best.safetensors

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

checkpoints/best_epoch_0.0145/best.pth

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

demo/video_3.mp4

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

evaluation.py

@@ -0,0 +1,179 @@
+import os
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+import face_alignment
+import lpips
+import pandas as pd
+
+from unet_acc import DenseMotion, UNetGenerator, warp_image
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+dense_motion = DenseMotion(kp_channels=68).to(device)
+generator = UNetGenerator(in_channels=4).to(device)
+
+ckpt = torch.load("checkpoints/best.pth", map_location=device)
+dense_motion.load_state_dict(ckpt["dense_motion"])
+generator.load_state_dict(ckpt["generator"])
+
+dense_motion.eval()
+generator.eval()
+
+lpips_fn = lpips.LPIPS(net="alex").to(device)
+
+fa = face_alignment.FaceAlignment(
+    face_alignment.LandmarksType.TWO_D,
+    device=device
+)
+
+
+# Metrics
+
+def landmark_distance(pred, gt):
+    pred = cv2.cvtColor(pred, cv2.COLOR_GRAY2RGB)
+    gt   = cv2.cvtColor(gt, cv2.COLOR_GRAY2RGB)
+
+    pl = fa.get_landmarks(pred)
+    gl = fa.get_landmarks(gt)
+    if pl is None or gl is None:
+        return None
+
+    pl, gl = pl[0], gl[0]
+    eye_dist = np.linalg.norm(gl[36] - gl[45]) + 1e-6
+    return np.mean(np.linalg.norm(pl - gl, axis=1)) / eye_dist
+
+
+def lpips_score(pred, gt):
+    pred = pred.repeat(1, 3, 1, 1)
+    gt   = gt.repeat(1, 3, 1, 1)
+    return lpips_fn(pred, gt).item()
+
+
+def l1_score(pred, gt):
+    return F.l1_loss(pred, gt).item()
+
+
+def temporal_jitter(frames):
+    diffs = []
+    for i in range(1, len(frames)):
+        diffs.append(torch.mean(torch.abs(frames[i] - frames[i - 1])).item())
+    return np.std(diffs), np.mean(diffs)
+
+
+LOCK_IDXS = list(range(36, 48)) + list(range(48, 68))
+
+def infer_no_warp(src):
+    B, _, H, W = src.shape
+    flow = torch.zeros(B, 2, H, W).to(device)
+    occ  = torch.ones(B, 1, H, W).to(device)
+    return torch.clamp(generator(torch.cat([src, flow, occ], 1)), 0, 1)
+
+
+def infer_warp(src, src_kp, drv_kp):
+    flow, occ = dense_motion(src_kp, drv_kp)
+    warped = warp_image(src, flow)
+    return torch.clamp(generator(torch.cat([warped, flow, occ], 1)), 0, 1)
+
+
+def infer_warp_lock(src, src_kp, drv_kp):
+    kp = src_kp.clone()
+    kp[:, LOCK_IDXS] = drv_kp[:, LOCK_IDXS]
+    flow, occ = dense_motion(src_kp, kp)
+    warped = warp_image(src, flow)
+    return torch.clamp(generator(torch.cat([warped, flow, occ], 1)), 0, 1)
+
+
+def infer_warp_lock_mask(src, src_kp, drv_kp, mask):
+    kp = src_kp.clone()
+    kp[:, LOCK_IDXS] = drv_kp[:, LOCK_IDXS]
+    flow, occ = dense_motion(src_kp, kp)
+    warped = warp_image(src, flow)
+    pred = generator(torch.cat([warped, flow, occ], 1))
+    return torch.clamp(pred * mask + src * (1 - mask), 0, 1)
+
+
+def evaluate_sequence(src, src_kp, drv_kps, gt_frames, mask, mode):
+    preds_torch = []
+    lmd, lp, l1 = [], [], []
+
+    with torch.no_grad():
+        for t, drv_kp in enumerate(drv_kps):
+            if mode == "no_warp":
+                pred = infer_no_warp(src)
+            elif mode == "warp":
+                pred = infer_warp(src, src_kp, drv_kp)
+            elif mode == "warp_lock":
+                pred = infer_warp_lock(src, src_kp, drv_kp)
+            elif mode == "warp_lock_mask":
+                pred = infer_warp_lock_mask(src, src_kp, drv_kp, mask)
+            else:
+                raise ValueError
+
+            gt = gt_frames[t]
+
+            pred_np = (pred.detach().cpu().squeeze().numpy() * 255).astype(np.uint8)
+            gt_np   = (gt.detach().cpu().squeeze().numpy() * 255).astype(np.uint8)
+
+            lm = landmark_distance(pred_np, gt_np)
+            if lm is not None:
+                lmd.append(lm)
+
+            lp.append(lpips_score(pred, gt))
+            l1.append(l1_score(pred, gt))
+            preds_torch.append(pred)
+
+    jit_std, _ = temporal_jitter(preds_torch)
+
+    return {
+        "LMD": np.mean(lmd) if len(lmd) > 0 else np.nan,
+        "LPIPS": np.mean(lp),
+        "Jitter": jit_std
+    }
+
+def run_all(src, src_kp, drv_kps, gt_frames, mask):
+    rows = []
+    for mode in ["no_warp", "warp", "warp_lock", "warp_lock_mask"]:
+        print(f"Evaluating {mode}")
+        res = evaluate_sequence(src, src_kp, drv_kps, gt_frames, mask, mode)
+        res["Method"] = mode
+        rows.append(res)
+
+    df = pd.DataFrame(rows)
+    df = df[["Method", "LMD", "LPIPS", "Jitter"]]
+    df.to_csv("ablation_results.csv", index=False)
+    print(df)
+if __name__ == "__main__":
+    src_img = cv2.imread(r"motion_transfer\new_dataset\test\dataset\87\frames\00000.jpg", cv2.IMREAD_GRAYSCALE)
+    src = torch.tensor(
+        src_img / 255.0,
+        dtype=torch.float32
+    ).unsqueeze(0).unsqueeze(0).to(device)
+
+    src_kp = torch.tensor(
+        np.load(r"motion_transfer\new_dataset\test\dataset\87\combined\00000.npy"),
+        dtype=torch.float32
+    ).permute(2, 0, 1).unsqueeze(0).to(device)
+
+    drv_kps = []
+    gt_frames = []
+
+    for f in sorted(os.listdir(r"motion_transfer\new_dataset\test\dataset\87\frames")):
+        gt = cv2.imread(os.path.join(r"motion_transfer\new_dataset\test\dataset\87\frames", f), cv2.IMREAD_GRAYSCALE)
+        gt_frames.append(
+            torch.tensor(
+                gt / 255.0,
+                dtype=torch.float32
+            ).unsqueeze(0).unsqueeze(0).to(device)
+        )
+
+        kp = torch.tensor(
+            np.load(os.path.join(r"motion_transfer\new_dataset\test\dataset\87\combined", f.replace(".jpg", ".npy"))),
+            dtype=torch.float32
+        ).permute(2, 0, 1).unsqueeze(0).to(device)
+        drv_kps.append(kp)
+
+    mask = torch.ones_like(src)
+    run_all(src, src_kp, drv_kps, gt_frames, mask)

model/unet_acc.py

@@ -0,0 +1,236 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+import numpy as np
+import cv2
+
+
+class SketchMotionDataset(Dataset):
+    def __init__(self, data_root, transform=None):
+        self.transform = transform
+        self.data = []
+
+        persons = sorted(os.listdir(data_root))
+        for person in persons:
+            person_dir = os.path.join(data_root, person)
+            source_frame = os.path.join(person_dir, f"{person}.jpg")
+            frames_dir = os.path.join(person_dir, "frames")
+            heatmap_dir = os.path.join(person_dir, "combined")
+
+            frame_files = sorted(os.listdir(frames_dir))
+            heatmap_files = sorted(os.listdir(heatmap_dir))
+
+            for f_file, h_file in zip(frame_files, heatmap_files):
+                self.data.append({
+                    "source_frame": source_frame,
+                    "driving_frame": os.path.join(frames_dir, f_file),
+                    "driving_heatmap": os.path.join(heatmap_dir, h_file),
+                    "source_heatmap": os.path.join(heatmap_dir, heatmap_files[0])  # first frame heatmap
+                })
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        src_img = Image.open(item["source_frame"]).convert("L")
+        drv_img = Image.open(item["driving_frame"]).convert("L")
+
+        src_img = torch.tensor(np.array(src_img)/255.0, dtype=torch.float32).unsqueeze(0)
+        drv_img = torch.tensor(np.array(drv_img)/255.0, dtype=torch.float32).unsqueeze(0)
+
+        src_kp = torch.tensor(np.load(item["source_heatmap"]), dtype=torch.float32).permute(2,0,1)
+        drv_kp = torch.tensor(np.load(item["driving_heatmap"]), dtype=torch.float32).permute(2,0,1)
+
+        return src_img, drv_img, src_kp, drv_kp
+
+
+class DenseMotion(nn.Module):
+    def __init__(self, kp_channels=68):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(kp_channels*2, 128, 7, padding=3),
+            nn.ReLU(),
+            nn.Conv2d(128, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 3, 3, padding=1)
+        )
+
+    def forward(self, src_kp, drv_kp):
+        x = torch.cat([src_kp, drv_kp], dim=1)
+        out = self.conv(x)
+        flow = out[:, :2, :, :]
+        occ = torch.sigmoid(out[:, 2:3, :, :])
+        return flow, occ
+
+class UNetGenerator(nn.Module):
+    def __init__(self, in_channels=4, out_channels=1):
+        super().__init__()
+
+        def conv_block(in_c, out_c):
+            return nn.Sequential(
+                nn.Conv2d(in_c, out_c, 3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(out_c, out_c, 3, padding=1),
+                nn.ReLU(inplace=True)
+            )
+
+        self.enc1 = conv_block(in_channels, 64)
+        self.pool1 = nn.MaxPool2d(2)
+        self.enc2 = conv_block(64, 128)
+        self.pool2 = nn.MaxPool2d(2)
+        self.bottleneck = conv_block(128, 256)
+        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.dec2 = conv_block(256, 128)
+        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.dec1 = conv_block(128, 64)
+        self.final = nn.Conv2d(64, out_channels, 1)
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool1(e1))
+        b = self.bottleneck(self.pool2(e2))
+        d2 = self.up2(b)
+        d2 = self.dec2(torch.cat([d2, e2], dim=1))
+        d1 = self.up1(d2)
+        d1 = self.dec1(torch.cat([d1, e1], dim=1))
+        return self.final(d1)
+
+def warp_image(img, flow):
+    B, C, H, W = img.shape
+    grid_y, grid_x = torch.meshgrid(torch.linspace(-1,1,H), torch.linspace(-1,1,W))
+    grid = torch.stack((grid_x, grid_y),2).unsqueeze(0).repeat(B,1,1,1).to(img.device)
+    flow_norm = flow.permute(0,2,3,1) / torch.tensor([W/2, H/2]).to(img.device)
+    warped = nn.functional.grid_sample(img, grid + flow_norm, align_corners=True)
+    return warped
+
+
+
+def save_checkpoint(state, is_best, checkpoint_dir="checkpoints", filename="last.pth", best_filename="best.pth"):
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    filepath = os.path.join(checkpoint_dir, filename)
+    torch.save(state, filepath)
+    if is_best:
+        bestpath = os.path.join(checkpoint_dir, best_filename)
+        torch.save(state, bestpath)
+        print(f"Saved new best checkpoint: {bestpath}")
+    else:
+        print(f"Saved checkpoint: {filepath}")
+
+def train(data_root, epochs=500, resume_checkpoint="checkpoints/last.pth"):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    dataset = SketchMotionDataset(data_root)
+    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
+
+    dense_motion = DenseMotion(kp_channels=68).to(device)
+    generator = UNetGenerator(in_channels=4).to(device)
+
+    optimizer = optim.Adam(list(dense_motion.parameters()) + list(generator.parameters()), lr=1e-4)
+    criterion = nn.L1Loss()
+
+    start_epoch = 0
+    best_loss = float("inf")
+
+
+    if os.path.exists(resume_checkpoint):
+        print(f"Resuming training from {resume_checkpoint}")
+        checkpoint = torch.load(resume_checkpoint, map_location=device)
+        dense_motion.load_state_dict(checkpoint["dense_motion"])
+        generator.load_state_dict(checkpoint["generator"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        start_epoch = checkpoint.get("epoch", 0)
+        best_loss = checkpoint.get("loss", float("inf"))
+        print(f"Resumed from epoch {start_epoch}, last loss = {best_loss:.4f}")
+    else:
+        print("Starting new training")
+
+    # Training  
+    for epoch in range(start_epoch, epochs):
+        epoch_loss = 0.0
+
+        for src_img, drv_img, src_kp, drv_kp in dataloader:
+            src_img, drv_img = src_img.to(device), drv_img.to(device)
+            src_kp, drv_kp = src_kp.to(device), drv_kp.to(device)
+
+            flow, occ = dense_motion(src_kp, drv_kp)
+            warped_src = warp_image(src_img, flow)
+
+            unet_input = torch.cat([warped_src, flow, occ], dim=1)
+            pred = generator(unet_input)
+
+            loss = criterion(pred, drv_img)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+
+        avg_loss = epoch_loss / len(dataloader)
+        print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f}")
+
+        # Save checkpoint
+        is_best = avg_loss < best_loss
+        best_loss = min(avg_loss, best_loss)
+
+        save_checkpoint({
+            "epoch": epoch + 1,
+            "dense_motion": dense_motion.state_dict(),
+            "generator": generator.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "loss": avg_loss,
+        }, is_best, checkpoint_dir="checkpoints")
+
+    print("Training completed")
+
+
+
+def generate_video(data_root, output_dir="outputs"):
+    os.makedirs(output_dir, exist_ok=True)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    dense_motion = DenseMotion(kp_channels=68).to(device)
+    generator = UNetGenerator(in_channels=4).to(device)
+
+    # Load best checkpoint
+    checkpoint = torch.load("checkpoints/best.pth", map_location=device)
+    dense_motion.load_state_dict(checkpoint["dense_motion"])
+    generator.load_state_dict(checkpoint["generator"])
+    dense_motion.eval()
+    generator.eval()
+
+    for person in sorted(os.listdir(data_root)):
+        person_dir = os.path.join(data_root, person)
+        source_frame = os.path.join(person_dir, f"{person}.jpg")
+        heatmap_dir = os.path.join(person_dir, "combined")
+
+        src_img = Image.open(source_frame).convert("L")
+        src_img = torch.tensor(np.array(src_img)/255.0, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)
+
+        src_kp = torch.tensor(np.load(os.path.join(heatmap_dir, sorted(os.listdir(heatmap_dir))[0])), dtype=torch.float32).permute(2,0,1).unsqueeze(0).to(device)
+
+        generated_frames = []
+        for h_file in sorted(os.listdir(heatmap_dir)):
+            drv_kp = torch.tensor(np.load(os.path.join(heatmap_dir, h_file)), dtype=torch.float32).permute(2,0,1).unsqueeze(0).to(device)
+
+            flow, occ = dense_motion(src_kp, drv_kp)
+            warped_src = warp_image(src_img, flow)
+            unet_input = torch.cat([warped_src, flow, occ], dim=1)
+            pred = generator(unet_input)
+            generated_frames.append(pred.detach().cpu().squeeze().numpy())
+
+        H, W = generated_frames[0].shape
+        out_path = os.path.join(output_dir, f"{person}_sketch.avi")
+        out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"XVID"), 15, (W,H), False)
+        for f in generated_frames:
+            out.write((f*255).astype(np.uint8))
+        out.release()
+        print(f"Video saved: {out_path}")
+
+if __name__ == "__main__":
+    data_root = "motion_transfer/dataset/"
+    train(data_root, epochs=500, resume_checkpoint="checkpoints/last.pth")  
+    generate_video(data_root, output_dir="outputs")

requirements.txt.txt

@@ -0,0 +1,22 @@
+diffusers==0.35.2
+face-alignment==1.4.1
+ffmpeg-python==0.2.0
+huggingface-hub==0.36.0
+imageio==2.37.0
+imageio-ffmpeg==0.4.7
+matplotlib==3.10.5
+mediapipe==0.10.21
+numpy==2.2.6
+opencv-contrib-python==4.11.0.86
+opencv-python==4.12.0.88
+safetensors==0.6.2
+scikit-image==0.25.2
+scikit-learn==1.7.1
+scipy==1.15.3
+sentencepiece==0.2.0
+torch==2.5.1+cu121
+torchaudio==2.5.1+cu121
+torchvision==0.20.1+cu121
+tqdm==4.64.1
+transformers==4.57.1
+urllib3==2.5.0

test_acc_upgrade.py

@@ -0,0 +1,210 @@
+import os
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+import mediapipe as mp
+from safetensors.torch import load_file
+
+from unet_acc import DenseMotion, UNetGenerator, warp_image
+
+
+
+# Face detector 
+
+mp_face_detection = mp.solutions.face_detection
+detector = mp_face_detection.FaceDetection(
+    model_selection=1,
+    min_detection_confidence=0.5
+)
+
+
+def crop_head_with_bg(img_rgb, target_size=256,
+                      margin_top=0.6, margin_sides=0.3, margin_bottom=0.4):
+    ih, iw, _ = img_rgb.shape
+    results = detector.process(img_rgb)
+    if not results.detections:
+        return None
+
+    det = results.detections[0]
+    bbox = det.location_data.relative_bounding_box
+
+    x1 = int(bbox.xmin * iw)
+    y1 = int(bbox.ymin * ih)
+    w = int(bbox.width * iw)
+    h = int(bbox.height * ih)
+
+    x1 = max(0, int(x1 - w * margin_sides))
+    x2 = min(iw, int(x1 + w * (1 + 2 * margin_sides)))
+    y1 = max(0, int(y1 - h * margin_top))
+    y2 = min(ih, int(y1 + h * (1 + margin_bottom + margin_top)))
+
+    cropped = img_rgb[y1:y2, x1:x2]
+    ch, cw = cropped.shape[:2]
+
+    scale = target_size / max(ch, cw)
+    new_w, new_h = int(cw * scale), int(ch * scale)
+    resized = cv2.resize(cropped, (new_w, new_h),
+                          interpolation=cv2.INTER_LANCZOS4)
+
+    blurred_bg = cv2.GaussianBlur(resized, (51, 51), 0)
+    background = cv2.resize(
+        blurred_bg, (target_size, target_size),
+        interpolation=cv2.INTER_AREA
+    )
+
+    y_off = (target_size - new_h) // 2
+    x_off = (target_size - new_w) // 2
+    background[y_off:y_off + new_h, x_off:x_off + new_w] = resized
+
+    return background
+
+
+
+# MediaPipe face mesh
+
+def get_mediapipe_keypoints(img_rgb):
+    h, w = img_rgb.shape[:2]
+    mp_face = mp.solutions.face_mesh
+    with mp_face.FaceMesh(static_image_mode=True, max_num_faces=1) as mesh:
+        res = mesh.process(img_rgb)
+        if not res.multi_face_landmarks:
+            raise RuntimeError("No face landmarks detected")
+        pts = [(p.x * w, p.y * h)
+               for p in res.multi_face_landmarks[0].landmark]
+    return np.array(pts, dtype=np.float32)
+
+
+
+
+
+def create_eye_mouth_mask(image_shape, keypoints):
+    H, W = image_shape
+    mask = np.zeros((H, W), dtype=np.uint8)
+
+    left_eye = [33, 133, 160, 159, 158, 157, 173]
+    right_eye = [362, 263, 387, 386, 385, 384, 398]
+    mouth_outer = list(range(61, 79))
+    mouth_inner = list(range(308, 325))
+
+    def fill(indices):
+        pts = keypoints[indices].astype(np.int32)
+        cv2.fillPoly(mask, [pts.reshape(-1, 1, 2)], 255)
+
+    fill(left_eye)
+    fill(right_eye)
+    fill(mouth_outer)
+    fill(mouth_inner)
+
+    mask = cv2.dilate(mask,
+                      cv2.getStructuringElement(
+                          cv2.MORPH_ELLIPSE, (5, 5)),
+                      1)
+    mask = cv2.GaussianBlur(mask, (7, 7), 2)
+    return mask.astype(np.float32) / 255.0
+
+
+
+# Inference
+
+def test_single_image(source_image_path, heatmap_dir, output_path):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    # Load model
+    ckpt = load_file("checkpoints/best.safetensors")
+
+    dense_motion = DenseMotion(kp_channels=68).to(device)
+    generator = UNetGenerator(in_channels=4).to(device)
+
+    dm_state = {k.replace("dense_motion.", ""): v
+                for k, v in ckpt.items()
+                if k.startswith("dense_motion.")}
+    gen_state = {k.replace("generator.", ""): v
+                 for k, v in ckpt.items()
+                 if k.startswith("generator.")}
+
+    dense_motion.load_state_dict(dm_state, strict=False)
+    generator.load_state_dict(gen_state, strict=False)
+    dense_motion.eval()
+    generator.eval()
+
+    print("Model loaded")
+
+    img_bgr = cv2.imread(source_image_path)
+    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+
+    cropped = crop_head_with_bg(img_rgb, target_size=256)
+    if cropped is None:
+        raise RuntimeError("Face not detected")
+
+    src_gray = cv2.cvtColor(cropped, cv2.COLOR_RGB2GRAY)
+    src_np = src_gray.astype(np.float32) / 255.0
+    H, W = src_np.shape
+
+    src_tensor = torch.tensor(src_np).unsqueeze(0).unsqueeze(0).to(device)
+
+    #  Mask 
+    keypoints = get_mediapipe_keypoints(cropped)
+    mask_np = create_eye_mouth_mask((H, W), keypoints)
+    mask_tensor = torch.tensor(mask_np).unsqueeze(0).unsqueeze(0).to(device)
+
+    # Heatmaps 
+    heatmap_files = sorted(
+        f for f in os.listdir(heatmap_dir) if f.endswith(".npy")
+    )
+    if not heatmap_files:
+        raise RuntimeError("No heatmaps found")
+
+    src_kp = torch.tensor(
+        np.load(os.path.join(heatmap_dir, heatmap_files[0])),
+        dtype=torch.float32
+    ).permute(2, 0, 1).unsqueeze(0).to(device)
+
+    out = cv2.VideoWriter(
+        output_path,
+        cv2.VideoWriter_fourcc(*"mp4v"),
+        15,
+        (W, H),
+        False
+    )
+
+    # Inference
+    with torch.no_grad():
+        for i, hfile in enumerate(heatmap_files):
+            drv_kp = torch.tensor(
+                np.load(os.path.join(heatmap_dir, hfile)),
+                dtype=torch.float32
+            ).permute(2, 0, 1).unsqueeze(0).to(device)
+
+            combined_kp = src_kp.clone()
+            for idx in list(range(36, 48)) + list(range(48, 68)):
+                combined_kp[:, idx] = drv_kp[:, idx]
+
+            flow, occ = dense_motion(src_kp, combined_kp)
+            warped = torch.clamp(warp_image(src_tensor, flow), 0, 1)
+
+            pred = torch.clamp(
+                generator(torch.cat([warped, flow, occ], dim=1)),
+                0, 1
+            )
+
+            final_frame = pred * mask_tensor + src_tensor * (1 - mask_tensor)
+            frame_np = (final_frame.cpu().squeeze().numpy() * 255).astype(np.uint8)
+
+            out.write(frame_np)
+
+            if i == 0:
+                cv2.imwrite("preview_streamlit_matched.png", frame_np)
+
+    out.release()
+    print(f"Output saved: {output_path}")
+
+
+
+if __name__ == "__main__":
+    test_single_image(
+        source_image_path="motion_transfer/test/87.jpg",
+        heatmap_dir="motion_transfer/test/combined/",
+        output_path="outputs/final_streamlit_matched.mp4"
+    )

utilis/Face_keypoints_generate.py

@@ -0,0 +1,146 @@
+import os
+import glob
+import cv2
+import numpy as np
+import torch
+import face_alignment
+from multiprocessing import Pool, cpu_count
+print("hello")
+
+sources_dir = r"motion_transfer\new_dataset\test\image"
+videos_dir  = r"motion_transfer\new_dataset\test\video"
+output_root = r"motion_transfer\new_dataset\test\dataset"
+
+num_workers = min(cpu_count(), 4)
+target_size = 256
+SIGMA = 2.0
+
+NUM_FACE_POINTS = 68  
+
+# Temporal smoothing 
+SMOOTH_ALPHA = 0.7
+
+# Models
+fa = face_alignment.FaceAlignment(
+    face_alignment.LandmarksType.TWO_D,
+    device='cuda' if torch.cuda.is_available() else 'cpu'
+)
+
+
+def resize_with_gradient_padding(img, target_size):
+    h, w = img.shape[:2]
+    scale = target_size / max(h, w)
+    img = cv2.resize(img, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_AREA)
+    h, w = img.shape[:2]
+    tl = np.mean(img[0:5, 0:5, :], axis=(0,1))
+    tr = np.mean(img[0:5, -5:, :], axis=(0,1))
+    grad = np.linspace(tl, tr, target_size)
+    bg = np.tile(grad, (target_size, 1, 1)).astype(np.uint8)
+    y0 = (target_size - h) // 2
+    x0 = (target_size - w) // 2
+    bg[y0:y0+h, x0:x0+w] = img
+    return bg
+
+def clip_xy(x, y, w, h):
+    return float(np.clip(x, 0, w - 1)), float(np.clip(y, 0, h - 1))
+
+def ema(prev_pts, curr_pts, alpha=SMOOTH_ALPHA):
+    if prev_pts is None or prev_pts.shape != curr_pts.shape:
+        return curr_pts
+    return alpha * curr_pts + (1.0 - alpha) * prev_pts
+
+def gaussian_heatmaps(points, H, W, sigma=2.0):
+    N = points.shape[0]
+    yy, xx = np.mgrid[0:H, 0:W].astype(np.float32)
+    heat = np.zeros((H, W, N), dtype=np.float32)
+    s2 = 2 * (sigma ** 2)
+    for i, (x, y) in enumerate(points):
+        d2 = (xx - x) ** 2 + (yy - y) ** 2
+        heat[..., i] = np.exp(-d2 / s2)
+    return heat
+
+
+def process_person(person_name):
+    source_path = None
+    for ext in [".png", ".jpg", ".jpeg"]:
+        p = os.path.join(sources_dir, person_name + ext)
+        if os.path.exists(p):
+            source_path = p
+            break
+    video_path  = os.path.join(videos_dir,  f"{person_name}.mp4")
+    if not (os.path.isfile(source_path) and os.path.isfile(video_path)):
+        print(f"Missing files for {person_name}")
+        return
+
+    print(f" Processing {person_name}...")
+    person_root  = os.path.join(output_root, person_name)
+    frames_dir = os.path.join(output_root, person_name, "frames")
+    combined_dir = os.path.join(output_root, person_name, "combined")
+    keypoints_preview_dir = os.path.join(output_root, person_name, "keypoints_preview")
+    os.makedirs(frames_dir, exist_ok=True)
+    os.makedirs(combined_dir, exist_ok=True)
+    os.makedirs(keypoints_preview_dir, exist_ok=True)
+
+    # Save resized
+    src_img = cv2.imread(source_path)
+    if src_img is not None:
+        src_ref = resize_with_gradient_padding(src_img, target_size)
+        cv2.imwrite(os.path.join(person_root, f"{person_name}.jpg"), src_ref)
+
+    cap = cv2.VideoCapture(video_path)
+    frame_idx = 0
+    prev_points = None
+
+    while True:
+        ok, frame_bgr = cap.read()
+        if not ok:
+            break
+
+        # resize
+        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+        frame_rgb = resize_with_gradient_padding(frame_rgb, target_size)
+        H, W = frame_rgb.shape[:2]
+
+        # 68 face landmarks
+        fa_out = fa.get_landmarks(frame_rgb)
+        if fa_out is None or len(fa_out) == 0 or fa_out[0].shape[0] < NUM_FACE_POINTS:
+            frame_idx += 1
+            continue
+        face68 = fa_out[0][:NUM_FACE_POINTS].astype(np.float32)
+
+        # Smooth for temporal stability
+        face68 = ema(prev_points, face68, alpha=SMOOTH_ALPHA)
+        prev_points = face68.copy()
+
+        # Heatmaps
+        heatmap = gaussian_heatmaps(face68, H, W, sigma=SIGMA)
+
+        # Save preview + .npy
+        vis = frame_rgb.copy()
+        for (x, y) in face68.astype(int):
+            cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)
+        vis_bgr = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
+
+        fname = f"{frame_idx:05d}"
+        frame_file = f"{frame_idx:05d}.jpg"
+        cv2.imwrite(os.path.join(frames_dir, frame_file), frame_rgb)
+        #cv2.imwrite(os.path.join(keypoints_preview_dir, f"{fname}.png"), vis_bgr)
+        np.save(os.path.join(combined_dir, f"{fname}.npy"), heatmap)
+
+        frame_idx += 1
+
+    cap.release()
+    print(f"Done: {person_name} | Frames: {frame_idx} | Points/frame: {NUM_FACE_POINTS}")
+
+
+if __name__ == "__main__":
+    image_files = []
+    image_files += glob.glob(os.path.join(sources_dir, "*.png"))
+    image_files += glob.glob(os.path.join(sources_dir, "*.jpg"))
+    image_files += glob.glob(os.path.join(sources_dir, "*.jpeg"))
+
+    people = [os.path.splitext(os.path.basename(p))[0] for p in image_files]
+
+    print("People found:", people)
+    with Pool(num_workers) as p:
+        p.map(process_person, people)

utilis/generate_heatmap.py

@@ -0,0 +1,68 @@
+import os
+import numpy as np
+import cv2
+
+
+single_heatmap_path = "motion_transfer/dataset_single/reference_heatmap/00000.npy"
+reference_heatmap_dir = "motion_transfer/dataset_single/reference_heatmap"  # contains 150 npy files
+output_dir = "motion_transfer/dataset_single/test_heatmap"
+preview_video = "motion_transfer/dataset_single/simulated_motion.mp4"
+
+os.makedirs(output_dir, exist_ok=True)
+
+single_heatmap = np.load(single_heatmap_path)
+H, W, C = single_heatmap.shape
+
+# Extract keypoints 
+def extract_keypoints(hmap):
+    kps = []
+    for i in range(hmap.shape[2]):
+        y, x = np.where(hmap[:, :, i] > 0)
+        if len(x) > 0:
+            kps.append([np.mean(x), np.mean(y)]) 
+        else:
+            kps.append([0, 0])
+    return np.array(kps, dtype=np.float32)
+
+single_kp = extract_keypoints(single_heatmap)
+
+# reference motion 
+ref_files = sorted([f for f in os.listdir(reference_heatmap_dir) if f.endswith(".npy")])
+ref_heatmaps = [np.load(os.path.join(reference_heatmap_dir, f)) for f in ref_files]
+ref_kp_list = [extract_keypoints(hm) for hm in ref_heatmaps]
+
+# Compute motion relative to first reference frame 
+ref_base_kp = ref_kp_list[0]
+motion_vectors = [kp - ref_base_kp for kp in ref_kp_list]  
+
+# Apply motion to single input keypoints 
+fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+video_writer = cv2.VideoWriter(preview_video, fourcc, 30, (W, H))
+
+def gaussian_heatmaps(points, H, W, sigma=2.0):
+    N = points.shape[0]
+    yy, xx = np.mgrid[0:H, 0:W].astype(np.float32)
+    heat = np.zeros((H, W, N), dtype=np.float32)
+    s2 = 2 * (sigma ** 2)
+    for i, (x, y) in enumerate(points):
+        d2 = (xx - x) ** 2 + (yy - y) ** 2
+        heat[..., i] = np.exp(-d2 / s2)
+    return heat
+
+for frame_idx, displacement in enumerate(motion_vectors):
+    moved_kp = single_kp + displacement
+
+    # Generate Gaussian heatmap for all points
+    new_heatmap = gaussian_heatmaps(moved_kp, H, W, sigma=2.0)
+    np.save(os.path.join(output_dir, f"{frame_idx:05d}.npy"), new_heatmap)
+
+    
+    frame_vis = np.zeros((H, W, 3), dtype=np.uint8)
+    for (x, y) in moved_kp.astype(int):
+        cv2.circle(frame_vis, (x, y), 2, (0, 255, 0), -1)
+    video_writer.write(frame_vis)
+
+video_writer.release()
+
+print(f"Simulated motion heatmaps saved in '{output_dir}'")
+print(f"Preview video saved as '{preview_video}'")

utilis/jitter.py

@@ -0,0 +1,50 @@
+import os
+import numpy as np
+
+heatmap_dir = "motion_transfer/dataset_single/test_heatmap"
+smoothed_dir = "motion_transfer/dataset_single/smoothed_heatmaps"
+os.makedirs(smoothed_dir, exist_ok=True)
+
+# Load existing heatmaps and extract keypoints
+heatmap_files = sorted([f for f in os.listdir(heatmap_dir) if f.endswith(".npy")])
+kp_list = [np.load(os.path.join(heatmap_dir, f)) for f in heatmap_files]
+
+# Extract keypoints from each heatmap
+def extract_kpoints(hm):
+    kps = []
+    for i in range(hm.shape[2]):
+        y, x = np.where(hm[:, :, i] > 0)
+        if len(x) > 0:
+            kps.append([np.mean(x), np.mean(y)])
+        else:
+            kps.append([0, 0])
+    return np.array(kps, dtype=np.float32)
+
+kp_list = [extract_kpoints(hm) for hm in kp_list]
+
+#Apply temporal smoothing
+def temporal_smoothing(kp_list, alpha=0.7):
+    smoothed = [kp_list[0].copy()]
+    for i in range(1, len(kp_list)):
+        new_kp = alpha * smoothed[-1] + (1 - alpha) * kp_list[i]
+        smoothed.append(new_kp)
+    return smoothed
+
+smoothed_kp_list = temporal_smoothing(kp_list, alpha=0.7)
+
+# Recompute heatmaps with same dimensions
+H, W, C = np.load(os.path.join(heatmap_dir, heatmap_files[0])).shape
+
+def gaussian_heatmaps(points, H, W, sigma=2.0):
+    N = points.shape[0]
+    yy, xx = np.mgrid[0:H, 0:W].astype(np.float32)
+    heat = np.zeros((H, W, N), dtype=np.float32)
+    s2 = 2 * (sigma ** 2)
+    for i, (x, y) in enumerate(points):
+        d2 = (xx - x) ** 2 + (yy - y) ** 2
+        heat[..., i] = np.exp(-d2 / s2)
+    return heat
+
+for idx, kp in enumerate(smoothed_kp_list):
+    new_hm = gaussian_heatmaps(kp, H, W, sigma=2.0)
+    np.save(os.path.join(smoothed_dir, heatmap_files[idx]), new_hm)