#!/usr/bin/env python3
"""DROID-SLAM keyframes + dense depths -> approximate orthographic color mosaic.
Not a PKS implementation, not a metric/survey true orthophoto. No ODM/MicMac.
Use --probe first. A CUDA DROID-SLAM installation and calibration are required.
"""
import argparse
from functools import lru_cache
import hashlib
import json
import logging
import math
from pathlib import Path
import shutil
import subprocess
import sys
import time
import numpy as np
LOG=logging.getLogger('droid_orthomosaic')
VIDEO_PATH='mainindexedvideo.mp4'
SCHEMA=1
def save_json(path,obj):
path=Path(path);tmp=path.with_name(path.name+'.tmp')
tmp.write_text(json.dumps(obj,indent=2),encoding='utf-8');tmp.replace(path)
def load_json(path):return json.loads(Path(path).read_text(encoding='utf-8'))
def digest(path):
h=hashlib.sha256()
with open(path,'rb') as f:
for b in iter(lambda:f.read(1024*1024),b''):h.update(b)
return h.hexdigest()
def positive(x,name,zero=False):
if not math.isfinite(x) or (x<0 if zero else x<=0):raise ValueError(name+' must be finite and '+('nonnegative' if zero else 'positive'))
def probe_video(path):
import cv2
cap=cv2.VideoCapture(str(path))
try:
if not cap.isOpened():raise RuntimeError('Cannot open local video: '+str(path))
ok,img=cap.read()
if not ok:raise RuntimeError('Cannot decode first video frame')
fps=float(cap.get(cv2.CAP_PROP_FPS));n=float(cap.get(cv2.CAP_PROP_FRAME_COUNT))
return dict(width=img.shape[1],height=img.shape[0],reported_fps=fps,reported_frame_count=n,
approximate_duration=n/fps if math.isfinite(fps) and fps>0 else None)
finally:cap.release()
def extract_frames(video_path,output_dir,seconds_interval=.125,*,max_frame_size=1600,start=0,end=None):
"""Chronological CFR sampling. Cache metadata is OUTSIDE image-only directory."""
import cv2
positive(seconds_interval,'sampling interval');positive(start,'start',True)
if max_frame_size<64:raise ValueError('max_frame_size must be >=64')
if end is not None and (not math.isfinite(end) or end<=start):raise ValueError('end must be finite and > start')
video=Path(video_path).expanduser().resolve()
if not video.is_file():raise FileNotFoundError(str(video))
folder=Path(output_dir).expanduser().resolve();folder.mkdir(parents=True,exist_ok=True)
manifest=folder.parent/(folder.name+'_manifest.json');st=video.stat()
spec=dict(schema=SCHEMA,video=str(video),size=st.st_size,mtime_ns=st.st_mtime_ns,
seconds_interval=seconds_interval,max_frame_size=max_frame_size,start=start,end=end)
if manifest.is_file():
old=load_json(manifest);rows=old.get('frames',[])
if old.get('spec')!=spec:raise RuntimeError('Extraction inputs/settings changed; choose new --work-dir/--frames-dir')
if rows and all((folder/r['name']).is_file() and digest(folder/r['name'])==r['sha256'] for r in rows):
return [str(folder/r['name']) for r in rows]
raise RuntimeError('Cached extracted frames missing/changed')
if any(folder.iterdir()):raise RuntimeError('Nonempty unmanaged frame directory; use --images-dir or choose a new directory')
cap=cv2.VideoCapture(str(video));rows=[]
try:
if not cap.isOpened():raise RuntimeError('Cannot open video')
fps=float(cap.get(cv2.CAP_PROP_FPS))
if not math.isfinite(fps) or fps<=0:raise RuntimeError('Invalid reported FPS; inspect/transcode video; FPS is not guessed')
reported=float(cap.get(cv2.CAP_PROP_FRAME_COUNT));count=int(reported) if math.isfinite(reported) and reported>0 else 0
index=0;next_time=start
while cap.grab():
t=index/fps
if end is not None and t>=end:break
if t+1e-9>=next_time:
ok,img=cap.retrieve()
if not ok or img is None:raise RuntimeError('Cannot decode frame '+str(index))
next_time=start+(math.floor((t-start)/seconds_interval+1e-9)+1)*seconds_interval
h,w=img.shape[:2];s=min(1,max_frame_size/max(h,w))
if s<1:img=cv2.resize(img,(max(1,round(w*s)),max(1,round(h*s))),interpolation=cv2.INTER_AREA)
name='frame_%010d.jpg'%index;path=folder/name
if not cv2.imwrite(str(path),img,[cv2.IMWRITE_JPEG_QUALITY,95]):raise RuntimeError('Cannot save '+str(path))
rows.append(dict(name=name,source_frame=index,seconds=t,sha256=digest(path)))
if len(rows)%100==0:LOG.info('Extracted %d SLAM input frames at %.1fs',len(rows),t)
index+=1
if end is None and count and index<count-max(5,round(fps)):raise RuntimeError('Video decoding stopped unexpectedly early')
if len(rows)<12:raise RuntimeError('Need at least 12 sampled frames; choose a longer segment or higher --input-fps')
save_json(manifest,dict(spec=spec,original_fps=fps,frames=rows))
return [str(folder/r['name']) for r in rows]
finally:cap.release()
def list_images(path):
folder=Path(path).expanduser().resolve()
if not folder.is_dir():raise FileNotFoundError(str(folder))
files=sorted((p for p in folder.iterdir() if p.suffix.lower() in ('.jpg','.jpeg','.png')),key=lambda p:p.name)
if len(files)<12:raise ValueError('Need at least 12 chronological input images')
return list(map(str,files))
def calibration_for_size(size,calib=None,calib_size=None,assume_hfov=None):
w,h=size
if calib is not None:
if assume_hfov is not None:raise ValueError('Choose --calib OR --assume-hfov')
if not calib_size:raise ValueError('--calib-size WIDTH HEIGHT is required with --calib')
v=np.asarray(np.loadtxt(calib),dtype=float).reshape(-1)
if len(v) not in (4,8,9,12) or not np.isfinite(v).all():raise ValueError('Calibration must contain fx fy cx cy and optional 4/5/8 OpenCV distortion coefficients')
cw,ch=calib_size
if cw<=0 or ch<=0 or abs((cw/ch)/(w/h)-1)>.005:raise ValueError('Calibration/source aspect ratios differ; use calibration for the actual video crop')
if v[0]<=0 or v[1]<=0 or not (0<=v[2]<cw and 0<=v[3]<ch):raise ValueError('Invalid focal length or principal point')
out=v.copy();out[[0,2]]*=w/cw;out[[1,3]]*=h/ch
return dict(values=out.tolist(),width=w,height=h,approximate=False,source_file=str(Path(calib).resolve()),calibration_size=[cw,ch])
if assume_hfov is None:raise ValueError('Supply --calib camera.txt --calib-size W H, or explicitly opt into --assume-hfov DEGREES')
if not math.isfinite(assume_hfov) or not 20<assume_hfov<150:raise ValueError('Assumed horizontal FOV must be 20..150 degrees')
f=w/(2*math.tan(math.radians(assume_hfov)/2))
LOG.warning('Using USER-ASSUMED intrinsics, NOT measured calibration; geometry can be wrong')
return dict(values=[f,f,(w-1)/2,(h-1)/2],width=w,height=h,approximate=True,assumed_hfov_degrees=assume_hfov)
def demo_shape(h,w):
factor=math.sqrt(384*512/(h*w));h1=int(h*factor);w1=int(w*factor)
return h1,w1,h1-h1%8,w1-w1%8
def run_logged(command,log_path,cwd=None):
LOG.info('Running: %s',subprocess.list2cmdline(command))
with open(log_path,'a',encoding='utf-8') as log:
log.write('\nCOMMAND: '+json.dumps(command)+'\n');log.flush()
p=subprocess.Popen(command,cwd=str(cwd) if cwd else None,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,
text=True,encoding='utf-8',errors='replace',bufsize=1)
try:
for line in p.stdout:print(line,end='',flush=True);log.write(line);log.flush()
code=p.wait()
except BaseException:
p.terminate()
try:p.wait(timeout=15)
except subprocess.TimeoutExpired:p.kill();p.wait()
raise
if code:raise RuntimeError('Subprocess failed (%s); see %s'%(code,log_path))
def check_droid(repo,weights,droid_python,work):
repo=Path(repo).expanduser().resolve();weights=Path(weights).expanduser().resolve()
if not (repo/'demo.py').is_file() or not weights.is_file():
raise ValueError('Provide an existing DROID-SLAM checkout and pretrained droid.pth')
work=Path(work).resolve();work.mkdir(parents=True,exist_ok=True)
check='import torch,lietorch,droid_backends; assert torch.cuda.is_available(), "CUDA is required"; print(torch.cuda.get_device_name(0)); print("GPU total bytes:",torch.cuda.get_device_properties(0).total_memory)'
run_logged([str(droid_python),'-c',check],work/'preflight.log',repo)
CONVERT_CODE=r"""
import sys,os,json,numpy as np,torch
from scipy.spatial.transform import Rotation
src,out=sys.argv[1:]
# This is ONLY the locally generated official-demo reconstruction, never an untrusted .pt.
try:data=torch.load(src,map_location='cpu',weights_only=True)
except TypeError:data=torch.load(src,map_location='cpu')
required={'tstamps','images','disps','poses','intrinsics'}
if not required.issubset(data):raise RuntimeError('Unsupported DROID reconstruction schema')
def a(k):return data[k].detach().cpu().numpy()
poses=a('poses').astype(np.float64);K=a('intrinsics').astype(np.float64)*8.0
D=a('disps').astype(np.float32);t=a('tstamps').astype(np.float64)
N=len(t)
if N<3 or poses.shape!=(N,7) or K.shape!=(N,4) or D.ndim!=3 or D.shape[0]!=N:raise RuntimeError('Invalid exported shapes')
if not np.isfinite(poses).all() or not np.isfinite(K).all() or not np.isfinite(t).all():raise RuntimeError('Nonfinite poses/calibration/timestamps')
if not np.allclose(t,np.round(t),atol=1e-4):raise RuntimeError('Expected integer demo image indexes, not physical timestamps')
# Stored pose is world-to-camera [tx,ty,tz,qx,qy,qz,qw]. Invert to camera-to-world.
Rcw=Rotation.from_quat(poses[:,3:]).as_matrix();Rwc=Rcw.transpose(0,2,1)
C=-np.einsum('nij,nj->ni',Rwc,poses[:,:3])
T=np.tile(np.eye(4),(N,1,1));T[:,:3,:3]=Rwc;T[:,:3,3]=C
if np.count_nonzero(np.isfinite(D)&(D>0))<100:raise RuntimeError('No usable upsampled inverse depths')
os.makedirs(out,exist_ok=True)
for name,array in [('camera_to_world',T),('intrinsics_full',K),('inverse_depth',D),('input_indexes',np.round(t).astype(np.int64))]:
np.save(os.path.join(out,name+'.npy'),array,allow_pickle=False)
print(json.dumps({'keyframes':N,'depth_shape':list(D.shape),'pose':'camera_to_world','intrinsics':'saved intrinsics multiplied by 8'}))
"""
def reconstruct(frame_paths,work_dir,*,repo,weights,droid_python=sys.executable,calibration,
buffer=512,filter_thresh=2.4,keyframe_thresh=4.0):
from PIL import Image
repo=Path(repo).expanduser().resolve();weights=Path(weights).expanduser().resolve();work=Path(work_dir).resolve()
if not (repo/'demo.py').is_file() or not weights.is_file():raise ValueError('Provide real --droid-repo checkout and --weights droid.pth')
demo=(repo/'demo.py').read_text(encoding='utf-8')
for flag in ('--reconstruction_path','--disable_vis','--buffer'):
if flag not in demo:raise RuntimeError('Unsupported upstream demo.py interface: missing '+flag)
if buffer<16:raise ValueError('--buffer must be >=16')
positive(filter_thresh,'filter_thresh');positive(keyframe_thresh,'keyframe_thresh')
paths=[Path(p).expanduser().resolve() for p in frame_paths]
if len(paths)<12:raise ValueError('At least 12 chronological frames required')
for p in paths:
with Image.open(p) as im:
if im.size!=(calibration['width'],calibration['height']):raise ValueError('Mixed image sizes or calibration/image-size mismatch')
sources=[]
for relative in ('demo.py','droid_slam/depth_video.py','droid_slam/droid.py','droid_slam/motion_filter.py'):
p=repo/relative
if not p.is_file():raise RuntimeError('Unsupported repo layout: '+relative)
sources.append(dict(path=relative,sha256=digest(p)))
git=subprocess.run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True,text=True) if shutil.which('git') else None
spec=dict(schema=SCHEMA,inputs=[dict(path=str(p),sha256=digest(p)) for p in paths],calibration=calibration,
weights_sha256=digest(weights),repo=str(repo),source_files=sources,git_commit=git.stdout.strip() if git and git.returncode==0 else None,
buffer=buffer,filter_thresh=filter_thresh,keyframe_thresh=keyframe_thresh,python=str(droid_python))
work.mkdir(parents=True,exist_ok=True);manifest=work/'droid_manifest.json';done=work/'droid_success.json';maps=work/'reconstruction'
if manifest.is_file():
if load_json(manifest)!=spec:raise RuntimeError('SLAM inputs/settings changed; use a NEW --work-dir')
else:
if (work/'slam_input').exists() or maps.exists():raise RuntimeError('Unmanaged reconstruction/staged images; choose a new work directory')
save_json(manifest,spec)
arrays=('camera_to_world.npy','intrinsics_full.npy','inverse_depth.npy','input_indexes.npy')
if done.is_file() and all((maps/n).is_file() for n in arrays):
hashes=load_json(done).get('hashes',{})
if all(digest(maps/n)==hashes.get(n) for n in arrays):return maps
raise RuntimeError('Cached SLAM arrays changed; use a new work directory')
# GPU/dependency errors surface here before launching the expensive demo.
check_droid(repo,weights,droid_python,work)
staged=work/'slam_input';staged.mkdir(exist_ok=True)
expected=[]
for i,p in enumerate(paths):
target=staged/('frame_%010d'%i+p.suffix.lower());expected.append(target.name)
if not target.is_file() or digest(target)!=spec['inputs'][i]['sha256']:shutil.copyfile(p,target)
if {p.name for p in staged.iterdir()}!=set(expected):raise RuntimeError('Unexpected files in SLAM input directory')
calib_path=work/'calibration_staged.txt';np.savetxt(calib_path,np.asarray(calibration['values'])[None],fmt='%.12g')
pt=work/'reconstruction.pt';pt.unlink(missing_ok=True)
command=[str(droid_python),str(repo/'demo.py'),'--imagedir',str(staged),'--calib',str(calib_path),
'--stride','1','--weights',str(weights),'--buffer',str(buffer),'--filter_thresh',str(filter_thresh),
'--keyframe_thresh',str(keyframe_thresh),'--disable_vis','--reconstruction_path',str(pt)]
t0=time.perf_counter();run_logged(command,work/'droid.log',repo)
if not pt.is_file():raise RuntimeError('DROID did not produce reconstruction.pt')
maps.mkdir(exist_ok=True)
run_logged([str(droid_python),'-c',CONVERT_CODE,str(pt),str(maps)],work/'conversion.log',repo)
data=MapData(maps,calibration,paths)
if len(data.ids)>=buffer:LOG.warning('Keyframe buffer is full; inspect tracking and rerun in a new work directory with a suitable buffer')
keyframes=[dict(keyframe=i,input_index=int(k),source_file=str(paths[k])) for i,k in enumerate(data.ids)]
save_json(work/'keyframes.json',keyframes)
save_json(done,dict(hashes={n:digest(maps/n) for n in arrays},slam_and_conversion_seconds=time.perf_counter()-t0,
retained_keyframes=len(data.ids),note='No claim that every input frame tracked successfully'))
return maps
class MapData:
def __init__(self,maps,calibration,paths):
maps=Path(maps);self.T=np.load(maps/'camera_to_world.npy',mmap_mode='r');self.K=np.load(maps/'intrinsics_full.npy',mmap_mode='r')
self.disp=np.load(maps/'inverse_depth.npy',mmap_mode='r');self.ids=np.load(maps/'input_indexes.npy',mmap_mode='r')
self.paths=list(map(Path,paths));self.calibration=calibration
n=len(self.ids)
if n<3 or self.T.shape!=(n,4,4) or self.K.shape!=(n,4) or self.disp.ndim!=3 or self.disp.shape[0]!=n:
raise ValueError('Unsupported reconstruction array shapes')
if len(set(map(int,self.ids)))!=n or np.any(np.diff(self.ids)<=0) or np.any(self.ids<0) or np.any(self.ids>=len(paths)):
raise ValueError('Invalid keyframe/input index mapping')
if not np.isfinite(self.T).all() or not np.isfinite(self.K).all() or np.any(self.K[:,:2]<=0):raise ValueError('Invalid pose/intrinsics arrays')
h,w=calibration['height'],calibration['width'];h1,w1,hd,wd=demo_shape(h,w)
if self.disp.shape[1:]!=(hd,wd):raise RuntimeError('DROID demo resizing has changed; this adapter will not guess the geometry')
c=np.asarray(calibration['values'][:4]);expected=c*np.array([w1/w,h1/h,w1/w,h1/h])
if not np.allclose(self.K,expected[None],rtol=2e-3,atol=.1):raise RuntimeError('Saved-intrinsics convention differs from inspected upstream (/8); refusing incorrect backprojection')
self.h1,self.w1=h1,w1;self.depth_h,self.depth_w=hd,wd
@lru_cache(maxsize=8)
def depth(self,i):
d=np.asarray(self.disp[i]);valid=np.isfinite(d)&(d>1e-8)
z=np.zeros_like(d,dtype=np.float32);z[valid]=1/d[valid]
if valid.any():
med=float(np.median(z[valid]));valid&=(z>=med*.2)&(z<=med*5)
z[~valid]=0
return z
def camera_frame(T):
views=np.asarray(T[:,:3,2],dtype=np.float64).copy();views/=np.linalg.norm(views,axis=1)[:,None]
mean=views.mean(axis=0);coherence=float(np.linalg.norm(mean))
if coherence<.7:raise RuntimeError('Reconstructed camera directions are not consistently nadir; inspect SLAM')
up=-mean/coherence;x=np.asarray(T[0,:3,0]).copy();x-=up*np.dot(x,up)
if np.linalg.norm(x)<1e-8:raise RuntimeError('Degenerate local horizontal axis')
x/=np.linalg.norm(x);y=np.cross(up,x);y/=np.linalg.norm(y)
return np.stack([x,y,up]),np.median(T[:,:3,3],axis=0),coherence
def backproject(u,v,z,K,T):
fx,fy,cx,cy=K
pc=np.column_stack(((u-cx)*z/fx,(v-cy)*z/fy,z))
return pc@T[:3,:3].T+T[:3,3]
def consistency(data,i,points,neighbors=2,tolerance=.06,min_support=1):
support=np.zeros(len(points),dtype=np.uint8)
for j in range(max(0,i-neighbors),min(len(data.ids),i+neighbors+1)):
if j==i:continue
pc=(points-data.T[j,:3,3])@data.T[j,:3,:3]
fx,fy,cx,cy=data.K[j];z=pc[:,2]
valid=np.isfinite(pc).all(axis=1)&(z>1e-8)
u=np.zeros(len(z));v=np.zeros(len(z));u[valid]=fx*pc[valid,0]/z[valid]+cx;v[valid]=fy*pc[valid,1]/z[valid]+cy
valid&=np.isfinite(u)&np.isfinite(v)&(u>=0)&(v>=0)&(u<data.depth_w-.5)&(v<data.depth_h-.5)
ix=np.zeros(len(z),dtype=np.int64);iy=np.zeros(len(z),dtype=np.int64)
ix[valid]=np.rint(u[valid]).astype(np.int64);iy[valid]=np.rint(v[valid]).astype(np.int64)
k=np.flatnonzero(valid)
other=data.depth(j)[iy[k],ix[k]]
good=(other>0)&(np.abs(other-z[k])<=tolerance*np.maximum(other,z[k]))
support[k[good]]+=1
return support>=min_support
def sample_geometry(data,i,stride=4,border=.03,neighbors=2,tolerance=.06,min_support=1):
h,w=data.depth_h,data.depth_w
yy,xx=np.mgrid[0:h:stride,0:w:stride];u=xx.ravel();v=yy.ravel();z=data.depth(i)[v,u]
ok=(z>0)&(u>=border*w)&(u<(1-border)*w)&(v>=border*h)&(v<(1-border)*h)
points=backproject(u[ok],v[ok],z[ok],data.K[i],data.T[i])
return points[consistency(data,i,points,neighbors,tolerance,min_support)]
def pixel_plan(lo,hi,long_edge,max_megapixels):
if not 256<=long_edge<=20000:raise ValueError('--long-edge must be 256..20000')
positive(max_megapixels,'max_megapixels')
span=np.asarray(hi)-lo
if not np.isfinite(span).all() or np.any(span<=0):raise ValueError('Degenerate XY bounds')
# Pixel CENTERS include both extent endpoints.
ppu=(long_edge-1)/max(span);w=int(math.ceil(span[0]*ppu))+1;h=int(math.ceil(span[1]*ppu))+1
if w*h>max_megapixels*1e6:raise ValueError('Raster exceeds --max-megapixels; lower --long-edge or explicitly increase budget')
return dict(width=w,height=h,pixels_per_model_unit=float(ppu),xy_min=list(map(float,lo)),xy_max=list(map(float,hi)))
def zbuffer_update(zbuf,scorebuf,rgbbuf,rows,cols,z,score,colors,height_tolerance):
"""Deterministic per-pixel top-surface selection, then view score near that surface."""
if not len(z):return
width=zbuf.shape[1];cell=rows.astype(np.int64)*width+cols
order=np.lexsort((-score,-z,cell));ordered=cell[order]
first=np.r_[True,ordered[1:]!=ordered[:-1]];take=order[first]
r,c=rows[take],cols[take];nz=z[take];ns=score[take]
old=zbuf[r,c];old_score=scorebuf[r,c]
replace=(nz>old+height_tolerance)|((np.abs(nz-old)<=height_tolerance)&(ns>old_score))
k=take[replace];r,c=rows[k],cols[k]
zbuf[r,c]=z[k];scorebuf[r,c]=score[k];rgbbuf[r,c]=colors[k]
def render_color_points(data,i,row_chunk=96,render_stride=1,border=.03):
import cv2
image=cv2.imread(str(data.paths[int(data.ids[i])]))
if image is None:raise RuntimeError('Cannot read keyframe source image')
cal=np.asarray(data.calibration['values']);K0=np.array([[cal[0],0,cal[2]],[0,cal[1],cal[3]],[0,0,1.]])
if len(cal)>4:image=cv2.undistort(image,K0,cal[4:])
h,w=image.shape[:2]
for y0 in range(0,h,row_chunk):
ys=np.arange(y0,min(h,y0+row_chunk),render_stride);xs=np.arange(0,w,render_stride)
yy,xx=np.meshgrid(ys,xs,indexing='ij')
u=xx.ravel();v=yy.ravel();ud=u*data.w1/w;vd=v*data.h1/h
inside=(ud>=border*data.depth_w)&(ud<(1-border)*data.depth_w)&(vd>=border*data.depth_h)&(vd<(1-border)*data.depth_h)
u,v,ud,vd=u[inside],v[inside],ud[inside],vd[inside]
# Bilinear interpolation of the estimated depth increases sampling density, NOT geometric information.
z=cv2.remap(data.depth(i),ud.astype(np.float32)[None],vd.astype(np.float32)[None],cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,borderValue=0).ravel()
near=data.depth(i)[np.clip(np.rint(vd).astype(int),0,data.depth_h-1),np.clip(np.rint(ud).astype(int),0,data.depth_w-1)]
ok=(z>0)&(near>0)&(np.abs(z-near)<=.1*np.maximum(z,near))
if not ok.any():continue
u,v,ud,vd,z=u[ok],v[ok],ud[ok],vd[ok],z[ok]
points=backproject(ud,vd,z,data.K[i],data.T[i])
fx,fy,cx,cy=data.K[i];score=1/(1+((ud-cx)/fx)**2+((vd-cy)/fy)**2)
# Crop source boundary and prefer image centers; color is from source, not low-res saved demo images.
score*=np.minimum.reduce([ud/data.depth_w,vd/data.depth_h,1-ud/data.depth_w,1-vd/data.depth_h])
yield points,score.astype(np.float32),image[v,u,::-1]
def export_raster(output_tif,output_image,rgb,zbuf,plan,axes,origin,fill_pixels=0,preview_max_size=None):
import rasterio
from rasterio.windows import Window
from rasterio.transform import from_origin
from PIL import Image
valid=np.isfinite(zbuf)
if not valid.any():raise RuntimeError('No valid orthographic pixels after visibility filtering')
if fill_pixels:
from scipy.ndimage import distance_transform_edt
distance,indices=distance_transform_edt(~valid,return_indices=True)
fill=(~valid)&(distance<=fill_pixels)
rgb[fill]=rgb[indices[0][fill],indices[1][fill]]
valid[fill]=True
del distance,indices
h,w=valid.shape;ppu=plan['pixels_per_model_unit'];lo=plan['xy_min'];hi=plan['xy_max']
# Local arbitrary-unit affine, NO CRS; pixel centers match the sampled local positions.
transform=from_origin(lo[0]-.5/ppu,hi[1]+.5/ppu,1/ppu,1/ppu)
temporary=output_tif.with_name(output_tif.stem+'.writing.tif')
with rasterio.open(temporary,'w',driver='GTiff',width=w,height=h,count=4,dtype='uint8',
transform=transform,compress='deflate',tiled=True,blockxsize=256,blockysize=256,BIGTIFF='IF_SAFER') as dst:
from rasterio.enums import ColorInterp
dst.colorinterp=(ColorInterp.red,ColorInterp.green,ColorInterp.blue,ColorInterp.alpha)
for row in range(0,h,256):
for col in range(0,w,256):
hh,ww=min(256,h-row),min(256,w-col);s=np.s_[row:row+hh,col:col+ww]
tile=np.asarray(rgb[s]).copy();tile[~valid[s]]=255
rgba=np.concatenate([tile,np.where(valid[s],255,0).astype(np.uint8)[...,None]],axis=2)
dst.write(np.moveaxis(rgba,2,0),window=Window(col,row,ww,hh))
dst.update_tags(description='DROID-SLAM approximate orthographic color mosaic; NOT survey-controlled',
model_units='arbitrary',orientation='camera-derived, not geographic north',
local_axes=json.dumps(axes.tolist()),local_origin=json.dumps(origin.tolist()))
with rasterio.open(temporary) as src:
if (src.width,src.height)!=(w,h):raise RuntimeError('Exported raster dimensions do not match render')
temporary.replace(output_tif)
viewing=np.asarray(rgb).copy();viewing[~valid]=255
image=Image.fromarray(viewing)
if preview_max_size:
positive(preview_max_size,'preview_max_size');image.thumbnail((preview_max_size,preview_max_size),Image.Resampling.LANCZOS)
tmp=output_image.with_name(output_image.stem+'.writing'+output_image.suffix)
image.save(tmp,**({'quality':95,'subsampling':0} if output_image.suffix.lower() in ('.jpg','.jpeg') else {}));tmp.replace(output_image)
return dict(width=w,height=h,viewing_width=image.width,viewing_height=image.height,valid_fraction=float(valid.mean()))
def render_orthomosaic(data,work,output,*,long_edge=6000,max_megapixels=64,neighbors=2,
tolerance=.06,min_support=1,render_stride=1,border=.03,fill_pixels=0,preview_max_size=None):
work=Path(work).resolve();output=Path(output).expanduser().resolve()
if output.suffix.lower() not in ('.jpg','.jpeg','.png'):raise ValueError('--output must be JPG or PNG')
if not 1<=neighbors<=10 or not 1<=min_support<=2*neighbors:raise ValueError('Invalid neighbor/support settings')
if not 0<tolerance<.5 or not 0<=border<.2:raise ValueError('Invalid depth tolerance/border')
if not 1<=render_stride<=8 or not 0<=fill_pixels<=5:raise ValueError('Invalid render stride/hole fill')
axes,origin,coherence=camera_frame(data.T);lo=np.full(2,np.inf);hi=np.full(2,-np.inf);supported=0;nonempty=0
for i in range(len(data.ids)):
p=sample_geometry(data,i,4,border,neighbors,tolerance,min_support)
if len(p):
local=(p-origin)@axes.T;lo=np.minimum(lo,local[:,:2].min(axis=0));hi=np.maximum(hi,local[:,:2].max(axis=0));supported+=len(p);nonempty+=1
if supported<100 or nonempty<3:raise RuntimeError('Insufficient cross-view-consistent geometry; do not lower thresholds blindly; inspect calibration/tracking')
# Small padding: bounds came from sampled geometry, not necessarily exact surface perimeter.
pad=(hi-lo)*.01;lo-=pad;hi+=pad
plan=pixel_plan(lo,hi,long_edge,max_megapixels);h,w=plan['height'],plan['width'];ppu=plan['pixels_per_model_unit']
LOG.info('Orthographic target %dx%d; %d retained keyframes, %d with sampled supported geometry',w,h,len(data.ids),nonempty)
raster_dir=work/'orthomosaic';raster_dir.mkdir(parents=True,exist_ok=True);output.parent.mkdir(parents=True,exist_ok=True)
# Disk-backed arrays prevent the entire z/color/score raster needing resident RAM simultaneously.
zbuf=np.memmap(raster_dir/'zbuffer.bin',dtype='float32',mode='w+',shape=(h,w));zbuf[:]=-np.inf
scorebuf=np.memmap(raster_dir/'score.bin',dtype='float32',mode='w+',shape=(h,w));scorebuf[:]=-np.inf
rgb=np.memmap(raster_dir/'rgb.bin',dtype='uint8',mode='w+',shape=(h,w,3));rgb[:]=255
scene_scale=float(np.median([np.median(data.depth(i)[data.depth(i)>0]) for i in range(len(data.ids)) if np.any(data.depth(i)>0)]))
t0=time.perf_counter();accepted=0
for i in range(len(data.ids)):
for points,score,color in render_color_points(data,i,render_stride=render_stride,border=border):
ok=consistency(data,i,points,neighbors,tolerance,min_support)
if not ok.any():continue
local=(points[ok]-origin)@axes.T;score=score[ok];color=color[ok]
col=np.rint((local[:,0]-lo[0])*ppu).astype(np.int64);row=np.rint((hi[1]-local[:,1])*ppu).astype(np.int64)
inside=(row>=0)&(col>=0)&(row<h)&(col<w)
if not inside.any():continue
zbuffer_update(zbuf,scorebuf,rgb,row[inside],col[inside],local[inside,2],score[inside],color[inside],scene_scale*.003)
accepted+=int(inside.sum())
LOG.info('Rendered keyframe %d/%d',i+1,len(data.ids))
zbuf.flush();scorebuf.flush();rgb.flush();tif=raster_dir/'orthomosaic.tif'
result=export_raster(tif,output,rgb,zbuf,plan,axes,origin,fill_pixels,preview_max_size)
save_json(work/'orthomosaic_report.json',dict(schema=SCHEMA,engine='DROID-SLAM + custom depth point-splat orthographic renderer',
retained_keyframes=len(data.ids),supported_sample_points=supported,keyframes_with_sampled_support=nonempty,
accepted_color_samples=accepted,plan=plan,result=result,view_coherence=coherence,
visibility=dict(neighbors=neighbors,min_other_view_support=min_support,relative_depth_tolerance=tolerance),
color_render_stride=render_stride,border_fraction=border,hole_fill_pixels=fill_pixels,
approximate_calibration=data.calibration.get('approximate',False),render_seconds=time.perf_counter()-t0,
local_axes=axes.tolist(),local_origin=origin.tolist(),
limitations='Approximate monocular depth-based orthographic mosaic; no mesh/surface optimization, radiometric equalization, '
'global seam optimization, independent loop detector, metric scale, georeferencing or accuracy guarantee. '
'Interpolated depth does not create detail; point splatting may leave holes and ghosts.'))
# Close mmap handles before removing scratch files (important on Windows).
del zbuf,scorebuf,rgb
for name in ('zbuffer.bin','score.bin','rgb.bin'):(raster_dir/name).unlink(missing_ok=True)
return tif
def read_preview(tif,max_size=None):
import rasterio
from rasterio.enums import Resampling
with rasterio.open(tif) as src:
scale=min(1,max_size/max(src.width,src.height)) if max_size else 1
h,w=max(1,round(src.height*scale)),max(1,round(src.width*scale))
rgb=src.read([1,2,3],out_shape=(3,h,w),resampling=Resampling.bilinear)
alpha=src.read(4,out_shape=(h,w),resampling=Resampling.nearest);rgb[:,alpha==0]=255
return np.ascontiguousarray(np.moveaxis(rgb,0,-1)[:,:,::-1])
def stitch_frames(frame_paths,*,work_dir='droid_work',repo,weights,calibration,long_edge=6000,
droid_python=sys.executable,output='orthomosaic.jpg',**render_options):
"""Legacy BGR-array contract with explicit new SLAM/calibration dependencies."""
maps=reconstruct(frame_paths,work_dir,repo=repo,weights=weights,droid_python=droid_python,calibration=calibration)
data=MapData(maps,calibration,frame_paths)
return read_preview(render_orthomosaic(data,work_dir,output,long_edge=long_edge,**render_options))
def main(argv=None):
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('video',nargs='?',default=VIDEO_PATH)
p.add_argument('--probe',action='store_true');p.add_argument('--images-dir')
p.add_argument('--check-droid',action='store_true',help='Test repo/weights/CUDA extensions without video extraction')
p.add_argument('--work-dir',default='droid_work');p.add_argument('--frames-dir')
p.add_argument('--droid-repo');p.add_argument('--weights');p.add_argument('--droid-python',default=sys.executable)
p.add_argument('--calib');p.add_argument('--calib-size',type=int,nargs=2,metavar=('WIDTH','HEIGHT'))
p.add_argument('--assume-hfov',type=float,help='Explicit approximate-intrinsics experiment, NOT measured camera calibration')
p.add_argument('--input-fps',type=float,default=8);p.add_argument('--interval',type=float,help='Overrides input-fps with seconds between samples')
p.add_argument('--max-frame-size',type=int,default=1600);p.add_argument('--start',type=float,default=0);p.add_argument('--end',type=float)
p.add_argument('--buffer',type=int,default=512);p.add_argument('--filter-thresh',type=float,default=2.4);p.add_argument('--keyframe-thresh',type=float,default=4)
p.add_argument('--render-only',action='store_true',help='Use saved maps/calibration; no CUDA/repo/weights needed')
p.add_argument('--long-edge',type=int,default=6000);p.add_argument('--max-megapixels',type=float,default=64)
p.add_argument('--neighbors',type=int,default=2);p.add_argument('--depth-tolerance',type=float,default=.06);p.add_argument('--min-support',type=int,default=1)
p.add_argument('--render-stride',type=int,default=1);p.add_argument('--border',type=float,default=.03)
p.add_argument('--hole-fill-pixels',type=int,default=0);p.add_argument('--preview-max-size',type=int)
p.add_argument('--output',default='orthomosaic.jpg')
a=p.parse_args(argv);logging.basicConfig(level=logging.INFO,format='%(levelname)s: %(message)s')
try:
if a.probe:print(json.dumps(probe_video(a.video),indent=2));return 0
if a.check_droid:
if not a.droid_repo or not a.weights:raise ValueError('--check-droid requires --droid-repo and --weights')
check_droid(a.droid_repo,a.weights,a.droid_python,a.work_dir);return 0
import rasterio # Missing export dependencies must fail BEFORE expensive reconstruction.
from PIL import Image
import cv2
if not 256<=a.long_edge<=20000:raise ValueError('--long-edge must be 256..20000')
positive(a.max_megapixels,'max_megapixels')
if a.preview_max_size is not None:positive(a.preview_max_size,'preview_max_size')
work=Path(a.work_dir).expanduser().resolve();work.mkdir(parents=True,exist_ok=True)
if a.render_only:
if a.calib or a.calib_size or a.assume_hfov is not None or a.images_dir or a.video!=VIDEO_PATH:
raise ValueError('--render-only uses saved sources/calibration; do not pass a new source or calibration')
spec=load_json(work/'droid_manifest.json');calibration=spec['calibration'];paths=[r['path'] for r in spec['inputs']]
if not all(Path(r['path']).is_file() and digest(r['path'])==r['sha256'] for r in spec['inputs']):raise RuntimeError('Source images changed/missing')
done=load_json(work/'droid_success.json')
if not all(digest(work/'reconstruction'/n)==h for n,h in done['hashes'].items()):raise RuntimeError('Cached reconstruction arrays changed')
maps=work/'reconstruction'
else:
if not a.droid_repo or not a.weights:raise ValueError('Supply --droid-repo and --weights, or use --render-only for an existing map')
if not a.calib and a.assume_hfov is None:raise ValueError('Measured intrinsics (--calib + --calib-size), or explicit --assume-hfov, are required')
if a.calib and (not a.calib_size or not Path(a.calib).is_file()):raise ValueError('Provide an existing --calib file AND --calib-size WIDTH HEIGHT')
if not (work/'droid_success.json').is_file():check_droid(a.droid_repo,a.weights,a.droid_python,work)
if a.images_dir:
if a.video!=VIDEO_PATH:raise ValueError('Choose positional video OR --images-dir')
paths=list_images(a.images_dir)
else:
positive(a.input_fps,'input_fps');interval=a.interval if a.interval is not None else 1/a.input_fps
paths=extract_frames(a.video,a.frames_dir or work/'frames',interval,max_frame_size=a.max_frame_size,start=a.start,end=a.end)
with Image.open(paths[0]) as im:size=im.size
calibration=calibration_for_size(size,a.calib,a.calib_size,a.assume_hfov)
maps=reconstruct(paths,work,repo=a.droid_repo,weights=a.weights,droid_python=a.droid_python,
calibration=calibration,buffer=a.buffer,filter_thresh=a.filter_thresh,keyframe_thresh=a.keyframe_thresh)
data=MapData(maps,calibration,paths)
tif=render_orthomosaic(data,work,a.output,long_edge=a.long_edge,max_megapixels=a.max_megapixels,
neighbors=a.neighbors,tolerance=a.depth_tolerance,min_support=a.min_support,render_stride=a.render_stride,
border=a.border,fill_pixels=a.hole_fill_pixels,preview_max_size=a.preview_max_size)
LOG.info('Saved full TIFF: %s',tif);LOG.info('Saved viewing image: %s',Path(a.output).resolve())
LOG.warning('Approximate depth-based orthographic mosaic, not a calibrated true orthophoto or a survey-controlled map.')
return 0
except KeyboardInterrupt:LOG.error('Interrupted; completed SLAM maps can be reused with --render-only');return 130
except Exception as exc:LOG.error('%s',exc);return 1
if __name__=='__main__':raise SystemExit(main())
