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