File size: 2,854 Bytes
acb2ce3 | 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 | #!/usr/bin/env python3
"""Predict real/fake for one image or a directory, without torchrun."""
import argparse
import json
import os
from pathlib import Path
import sys
from ltf_final import LTFDetector
from ltf_final.detector import ROOT, BRANCHES
from utils.tile_dataset_srm import IMG_EXTENSIONS
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--input', type=Path, required=True, help='Image file or directory, scanned recursively')
parser.add_argument('--output', type=Path, help='JSONL predictions; stdout is used when omitted')
parser.add_argument('--config', type=Path, default=ROOT / 'configs/eval.yaml')
parser.add_argument('--weights', type=Path, help='Override the bundled epoch_20.pth')
parser.add_argument('--device', default='auto', help='auto, cpu, cuda:0, etc.')
parser.add_argument('--branch', choices=BRANCHES, default='fusion')
parser.add_argument('--threshold', type=float, default=.5)
parser.add_argument('--chunk-size', type=int, default=256, help='CLIP views per forward chunk; tile coverage is preserved')
parser.add_argument('--threads', type=int, default=4, help='PyTorch CPU thread count')
args = parser.parse_args()
if args.threads < 1:
parser.error('--threads must be positive')
if not 0 <= args.threshold <= 1:
parser.error('--threshold must lie in [0,1]')
source = args.input.expanduser().resolve()
if source.is_file():
paths = [source]
elif source.is_dir():
paths = sorted(p for p in source.rglob('*') if p.is_file() and p.suffix.lower() in IMG_EXTENSIONS)
else:
parser.error(f'Input does not exist: {source}')
if not paths:
parser.error('No images found')
import torch
torch.set_num_threads(args.threads)
detector = LTFDetector(args.config, args.weights.resolve() if args.weights else None,
args.device, args.chunk_size)
stream, temporary = sys.stdout, None
if args.output:
output = args.output.expanduser().resolve()
output.parent.mkdir(parents=True, exist_ok=True)
temporary = output.with_name(output.name + f'.{os.getpid()}.tmp')
stream = temporary.open('w', encoding='utf-8')
try:
for path in paths:
result = detector.predict(path, args.branch, args.threshold)
stream.write(json.dumps(result, ensure_ascii=False, allow_nan=False) + '\n')
stream.flush()
if temporary:
stream.close()
os.replace(temporary, output)
except BaseException:
if temporary:
stream.close()
temporary.unlink(missing_ok=True)
raise
print(f'Predicted {len(paths)} image(s); branch={args.branch}; threshold={args.threshold}', file=sys.stderr)
if __name__ == '__main__':
main()
|