File size: 2,957 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 | """Dự đoán nhãn cho 1 ảnh bằng DINOv3 backbone + linear probe.
Linear probe được train lại trên đặc trưng đã cache (features.npz) — cùng split 80/20
như script đánh giá, nên kết quả dự đoán là hợp lệ (ảnh không nằm trong tập train).
"""
import argparse
import os
import sys
import numpy as np
import torch
from PIL import Image
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from torchvision import transforms
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from src.models.dinov3_vit import load_dinov3 # noqa: E402
IMG_SIZE = 256
TRANSFORM = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def train_probe(features_path: str, seed: int = 42, test_size: float = 0.2):
"""Train linear probe trên train split của features đã cache."""
data = np.load(features_path, allow_pickle=True)
X, y = data["features"], data["labels"]
X_train, _, y_train, _ = train_test_split(
X, y, test_size=test_size, stratify=y, random_state=seed
)
clf = make_pipeline(
StandardScaler(), LogisticRegression(max_iter=2000, random_state=seed)
)
clf.fit(X_train, y_train)
return clf
def extract_feature(model, img_path: str, device: str) -> np.ndarray:
img = TRANSFORM(Image.open(img_path).convert("RGB")).unsqueeze(0).to(device)
model.eval()
with torch.no_grad():
feat = model(img)
return feat.cpu().numpy()[0]
def main():
parser = argparse.ArgumentParser(description="Dự đoán nhãn 1 ảnh")
parser.add_argument("--image", required=True, help="Đường dẫn ảnh cần dự đoán")
parser.add_argument("--model", default="models/model.safetensors")
parser.add_argument("--features", default="outputs/results/features.npz")
parser.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"])
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"
# Train linear probe (trên train split — ảnh cần test nằm ngoài split này)
clf = train_probe(args.features)
model = load_dinov3(args.model, img_size=IMG_SIZE).to(device)
feat = extract_feature(model, args.image, device)
prob = clf.predict_proba(feat.reshape(1, -1))[0]
p_fake = float(prob[1]) # class 1 = FAKE
label = "FAKE" if p_fake >= 0.5 else "REAL"
print(f"\nẢnh: {args.image}")
print(f" → Kết luận: **{label}** (xác suất fake = {p_fake:.1%})")
print(f" real: {prob[0]:.1%} | fake: {prob[1]:.1%}")
if __name__ == "__main__":
main()
|