| |
| """ |
| Extract DepthAnythingV2 depth maps for cuneiform crops. |
| Saves as uint8 grayscale PNGs alongside each crop under depth/ subfolder. |
| |
| HF model: depth-anything/Depth-Anything-V2-Small-hf (fastest). |
| """ |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
|
|
| def load_depth_model(device='cuda'): |
| from transformers import AutoImageProcessor, AutoModelForDepthEstimation |
| proc = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf") |
| model = AutoModelForDepthEstimation.from_pretrained( |
| "depth-anything/Depth-Anything-V2-Small-hf", torch_dtype=torch.float16 |
| ).to(device).eval() |
| return proc, model |
|
|
|
|
| def depth_of(img, proc, model, device='cuda'): |
| inputs = proc(images=img, return_tensors="pt").to(device) |
| inputs["pixel_values"] = inputs["pixel_values"].to(torch.float16) |
| with torch.no_grad(): |
| out = model(**inputs) |
| pred = out.predicted_depth |
| pred = torch.nn.functional.interpolate( |
| pred.unsqueeze(1), size=img.size[::-1], mode="bicubic", align_corners=False |
| )[0, 0] |
| arr = pred.float().cpu().numpy() |
| |
| arr = arr - arr.min() |
| if arr.max() > 0: |
| arr = arr / arr.max() * 255 |
| return arr.astype(np.uint8) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--manifest', required=True) |
| ap.add_argument('--output-dir', default='/arf/scratch/stakan/hitit-proje/hitit_ocr/data/classification/depth') |
| ap.add_argument('--batch', type=int, default=1) |
| ap.add_argument('--limit', type=int, default=0) |
| args = ap.parse_args() |
|
|
| proc, model = load_depth_model() |
| out_dir = Path(args.output_dir); out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| seen = set() |
| n = 0 |
| with open(args.manifest) as f: |
| for line in f: |
| r = json.loads(line) |
| p = r.get('path') |
| if not p or p in seen: |
| continue |
| seen.add(p) |
| |
| try: |
| rel = Path(p).relative_to('/arf/scratch/stakan/hitit-proje/hitit_ocr/data/classification/all') |
| out_p = out_dir / rel |
| except ValueError: |
| |
| parts = Path(p).parts |
| key = '_'.join(parts[-3:]) |
| out_p = out_dir / '_external' / key |
| out_p = out_p.with_suffix('.png') |
| if out_p.exists(): |
| n += 1 |
| continue |
| out_p.parent.mkdir(parents=True, exist_ok=True) |
| try: |
| img = Image.open(p).convert('RGB') |
| d = depth_of(img, proc, model) |
| Image.fromarray(d).save(out_p) |
| except Exception as e: |
| print(f"fail {p}: {e}") |
| continue |
| n += 1 |
| if n % 500 == 0: |
| print(f" {n} depth maps") |
| if args.limit and n >= args.limit: |
| break |
| print(f"DONE: {n} depth maps → {out_dir}") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
