File size: 4,413 Bytes
e23cf4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """Trích xuất đặc trưng CLS token từ DINOv3 backbone cho toàn bộ ảnh trong data/test.
Output: outputs/results/features.npz gồm:
features (N, 384) — đặc trưng CLS token
labels (N,) — 0 = real, 1 = fake
paths (N,) — đường dẫn ảnh
"""
import argparse
import glob
import os
import sys
import numpy as np
import torch
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from src.models.dinov3_vit import load_dinov3 # noqa: E402
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
IMG_SIZE = 256
TRANSFORM = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])
LABELS = {"real": 0, "fake": 1}
def collect_images(data_dir: str):
"""Gom tất cả ảnh trong data_dir/{real,fake}. Trả về list (path, label)."""
items = []
for cls, label in LABELS.items():
folder = os.path.join(data_dir, cls)
if not os.path.isdir(folder):
print(f"[Cảnh báo] Không tìm thấy thư mục: {folder}")
continue
for ext in ("*.jpg", "*.jpeg", "*.png"):
items.extend((p, label) for p in glob.glob(os.path.join(folder, ext)))
return items
def collect_from_csv(csv_path: str):
"""Đọc manifest CSV (path,label). Trả về list (path, label)."""
import csv
items = []
with open(csv_path, newline="") as f:
reader = csv.reader(f)
header = next(reader, None)
if header and header != ["path", "label"]:
print(f"[Cảnh báo] Header CSV không chuẩn: {header}")
for row in reader:
if len(row) < 2:
continue
p, lbl = row[0], int(row[1])
if not os.path.exists(p):
print(f"[Cảnh báo] Thiếu file: {p}")
continue
items.append((p, lbl))
return items
def extract_features(model, items, device, batch_size=32):
model.to(device).eval()
feats, labels, paths = [], [], []
for i in tqdm(range(0, len(items), batch_size), desc="Extracting features", unit="batch"):
batch = items[i : i + batch_size]
imgs = [TRANSFORM(Image.open(p).convert("RGB")) for p, _ in batch]
x = torch.stack(imgs).to(device)
with torch.no_grad():
f = model(x)
feats.append(f.cpu().numpy())
labels.extend(lbl for _, lbl in batch)
paths.extend(p for p, _ in batch)
return np.vstack(feats), np.array(labels), np.array(paths, dtype=object)
def main():
parser = argparse.ArgumentParser(description="Trích xuất đặc trưng DINOv3")
parser.add_argument("--data-dir", default="data/test", help="Thư mục chứa {real, fake}")
parser.add_argument("--csv", default=None, help="Manifest CSV (path,label). Nếu có, bỏ qua --data-dir")
parser.add_argument("--model", default="models/model.safetensors")
parser.add_argument("--output", default="outputs/results/features.npz")
parser.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"])
parser.add_argument("--batch-size", type=int, default=32)
parser.add_argument("--act", default="silu", choices=["silu", "gelu"])
args = parser.parse_args()
if args.device == "auto":
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
else:
device = args.device
print(f"Device: {device}")
print(f"Load model: {args.model}")
model = load_dinov3(args.model, act=args.act, img_size=IMG_SIZE)
items = collect_from_csv(args.csv) if args.csv else collect_images(args.data_dir)
print(f"Tổng ảnh: {len(items)} (real={sum(1 for _, l in items if l == 0)}, fake={sum(1 for _, l in items if l == 1)})")
if not items:
sys.exit("Không có ảnh nào — kiểm tra lại --data-dir / --csv")
feats, labels, paths = extract_features(model, items, device, args.batch_size)
os.makedirs(os.path.dirname(args.output), exist_ok=True)
np.savez(args.output, features=feats, labels=labels, paths=paths)
print(f"Đã lưu {args.output}: features={feats.shape}, labels={labels.shape}")
if __name__ == "__main__":
main()
|