| """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 |
|
|
| 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" |
|
|
| |
| 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]) |
| 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() |
|
|
