diff --git a/.gitignore b/.gitignore index 4fbe3ef..59637a3 100644 --- a/.gitignore +++ b/.gitignore @@ -5,9 +5,8 @@ *.ckpt *.safetensors -# Пропсы конвейера NVIDIA (274 МБ): не наши и скачиваются штатным скриптом -assets/conveyors/ -# восстановить: python3 scripts/fetch_assets.py +# Наборы кадров прежних прогонов CV (1.26 ГБ) - выход, не исходники +cv/flow_*/ # Выход прогонов - воспроизводится, в истории не нужен control_test/captures/ diff --git a/assets/conveyors/ConveyorBelt_A03.usd b/assets/conveyors/ConveyorBelt_A03.usd new file mode 100644 index 0000000..6191bef Binary files /dev/null and b/assets/conveyors/ConveyorBelt_A03.usd differ diff --git a/assets/conveyors/ConveyorBelt_A06.usd b/assets/conveyors/ConveyorBelt_A06.usd new file mode 100644 index 0000000..76c0446 Binary files /dev/null and b/assets/conveyors/ConveyorBelt_A06.usd differ diff --git a/assets/conveyors/ConveyorBelt_A09.usd b/assets/conveyors/ConveyorBelt_A09.usd new file mode 100644 index 0000000..437405a Binary files /dev/null and b/assets/conveyors/ConveyorBelt_A09.usd differ diff --git a/assets/conveyors/ConveyorBelt_A21.usd b/assets/conveyors/ConveyorBelt_A21.usd new file mode 100644 index 0000000..5155ac5 Binary files /dev/null and b/assets/conveyors/ConveyorBelt_A21.usd differ diff --git a/assets/conveyors/ConveyorBelt_A24.usd b/assets/conveyors/ConveyorBelt_A24.usd new file mode 100644 index 0000000..825bbf7 Binary files /dev/null and b/assets/conveyors/ConveyorBelt_A24.usd differ diff --git a/assets/conveyors/Material Library/Metal/Misc/Metal_Rough_A.mdl b/assets/conveyors/Material Library/Metal/Misc/Metal_Rough_A.mdl new file mode 100644 index 0000000..de36791 --- /dev/null +++ b/assets/conveyors/Material Library/Metal/Misc/Metal_Rough_A.mdl @@ -0,0 +1,16 @@ +mdl 1.4; + +using ::OmniPBR import OmniPBR; +import ::tex::gamma_mode; +import ::state::normal; +import ::base::mono_mode; + +export material Metal_Rough_A(*) + = OmniPBR( + diffuse_texture: texture_2d("../Textures/T_Metal_Rough_A_Albedo.png", ::tex::gamma_srgb), + enable_ORM_texture: true, + metallic_texture_influence: 1.0, + normalmap_texture: texture_2d("../Textures/T_Metal_Rough_A_Normal.png", ::tex::gamma_linear), + ORM_texture: texture_2d("../Textures/T_Metal_Rough_A_ORM.png", ::tex::gamma_linear), + reflection_roughness_texture_influence: 1.0, + texture_scale: float2(2, 2)); diff --git a/assets/conveyors/Material Library/Metal/Painted/MetalPainted_White_Glossy_A.mdl b/assets/conveyors/Material Library/Metal/Painted/MetalPainted_White_Glossy_A.mdl new file mode 100644 index 0000000..8bb06bf --- /dev/null +++ b/assets/conveyors/Material Library/Metal/Painted/MetalPainted_White_Glossy_A.mdl @@ -0,0 +1,15 @@ +mdl 1.4; + +using ::OmniPBR import OmniPBR; +import ::tex::gamma_mode; +import ::state::normal; +import ::base::mono_mode; + +export material MetalPainted_White_Glossy_A(*) + = OmniPBR( + diffuse_texture: texture_2d("./T_MetalPainted_White_Glossy_A_Albedo.png", ::tex::gamma_srgb), + enable_ORM_texture: true, + metallic_texture_influence: 1.0, + normalmap_texture: texture_2d("./T_MetalPainted_White_Glossy_A_Normal.png", ::tex::gamma_linear), + ORM_texture: texture_2d("./T_MetalPainted_White_Glossy_A_ORM.png", ::tex::gamma_linear), + reflection_roughness_texture_influence: 1.0); diff --git a/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_Albedo.png b/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_Albedo.png new file mode 100644 index 0000000..04b161b Binary files /dev/null and b/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_Albedo.png differ diff --git a/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_Normal.png b/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_Normal.png new file mode 100644 index 0000000..6b49de4 Binary files /dev/null and b/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_Normal.png differ diff --git a/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_ORM.png b/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_ORM.png new file mode 100644 index 0000000..30344e1 Binary files /dev/null and b/assets/conveyors/Material Library/Metal/Painted/T_MetalPainted_White_Glossy_A_ORM.png differ diff --git a/assets/conveyors/Material Library/Metal/Steel/Steel_A.mdl b/assets/conveyors/Material Library/Metal/Steel/Steel_A.mdl new file mode 100644 index 0000000..33b9dbe --- /dev/null +++ b/assets/conveyors/Material Library/Metal/Steel/Steel_A.mdl @@ -0,0 +1,15 @@ +mdl 1.4; + +using ::OmniPBR import OmniPBR; +import ::tex::gamma_mode; +import ::state::normal; +import ::base::mono_mode; + +export material Steel_A(*) + = OmniPBR( + diffuse_texture: texture_2d("../Textures/T_Steel_A1_Albedo.png", ::tex::gamma_srgb), + enable_ORM_texture: true, + metallic_texture_influence: 1.0, + normalmap_texture: texture_2d("../Textures/T_Steel_A1_Normal.png", ::tex::gamma_linear), + ORM_texture: texture_2d("../Textures/T_Steel_A1_ORM.png", ::tex::gamma_linear), + reflection_roughness_texture_influence: 1.0); diff --git a/assets/conveyors/Material Library/Plastic/Misc/Plastic_Orange_A.mdl b/assets/conveyors/Material Library/Plastic/Misc/Plastic_Orange_A.mdl new file mode 100644 index 0000000..f224e75 --- /dev/null +++ b/assets/conveyors/Material Library/Plastic/Misc/Plastic_Orange_A.mdl @@ -0,0 +1,52 @@ +mdl 1.4; + +using ::OmniPBR import OmniPBR; +import ::tex::gamma_mode; +import ::state::normal; +import ::base::mono_mode; + +export material Plastic_Orange_A(*) + = OmniPBR( + diffuse_color_constant: color(0.2, 0.2, 0.2), + diffuse_texture: texture_2d("../Textures/T_Plastic_Orange_A_Albedo.png", ::tex::gamma_srgb), + albedo_desaturation: 0.0, + albedo_add: 0.0, + albedo_brightness: 1.0, + diffuse_tint: color(1, 1, 1), + reflection_roughness_constant: 0.5, + reflection_roughness_texture_influence: 1.0, + reflectionroughness_texture: texture_2d(), + metallic_constant: 0.0, + metallic_texture_influence: 1.0, + metallic_texture: texture_2d(), + specular_level: 0.5, + enable_ORM_texture: true, + ORM_texture: texture_2d("../Textures/T_Plastic_Orange_A_ORM.png", ::tex::gamma_linear), + ao_to_diffuse: 0.0, + ao_texture: texture_2d(), + enable_emission: false, + emissive_color: color(1, 0.1, 0.1), + emissive_color_texture: texture_2d(), + emissive_mask_texture: texture_2d(), + emissive_intensity: 40.0, + enable_opacity: false, + enable_opacity_texture: false, + opacity_constant: 1.0, + opacity_texture: texture_2d(), + opacity_mode: ::base::mono_average, + opacity_threshold: 0.0, + bump_factor: 1.0, + normalmap_texture: texture_2d("../Textures/T_Plastic_Orange_A_Normal.png", ::tex::gamma_linear), + detail_bump_factor: 0.30000001192092896, + detail_normalmap_texture: texture_2d(), + flip_tangent_u: false, + flip_tangent_v: true, + project_uvw: false, + world_or_object: false, + uv_space_index: 0, + texture_translate: float2(0, 0), + texture_rotate: 0.0, + texture_scale: float2(3, 3), + detail_texture_translate: float2(0, 0), + detail_texture_rotate: 0.0, + detail_texture_scale: float2(1, 1)); diff --git a/assets/conveyors/Material Library/Plastic/Misc/Plastic_Red_A.mdl b/assets/conveyors/Material Library/Plastic/Misc/Plastic_Red_A.mdl new file mode 100644 index 0000000..91a7ea2 --- /dev/null +++ b/assets/conveyors/Material Library/Plastic/Misc/Plastic_Red_A.mdl @@ -0,0 +1,20 @@ +mdl 1.4; + +using ::OmniPBR import OmniPBR; +import ::tex::gamma_mode; +import ::state::normal; +import ::base::mono_mode; + +export material Plastic_Red_A(*) + = OmniPBR( + albedo_add: 0.4099999964237213, + albedo_brightness: 0.9099999666213989, + diffuse_texture: texture_2d("../Textures/T_Plastic_Red_A_Albedo.png", ::tex::gamma_srgb), + diffuse_tint: color(0.2680115, 0.053293347, 0.053293347), + enable_ORM_texture: true, + metallic_constant: 0.0, + metallic_texture_influence: 1.0, + normalmap_texture: texture_2d("../Textures/T_Plastic_Red_A_Normal.png", ::tex::gamma_linear), + ORM_texture: texture_2d("../Textures/T_Plastic_Red_A_ORM.png", ::tex::gamma_linear), + reflection_roughness_texture_influence: 1.0, + texture_scale: float2(3, 3)); diff --git a/assets/conveyors/Material Library/Plastic/Misc/Plastic_Rough_Black_A.mdl b/assets/conveyors/Material Library/Plastic/Misc/Plastic_Rough_Black_A.mdl new file mode 100644 index 0000000..d5a824d --- /dev/null +++ b/assets/conveyors/Material Library/Plastic/Misc/Plastic_Rough_Black_A.mdl @@ -0,0 +1,52 @@ +mdl 1.4; + +using ::OmniPBR import OmniPBR; +import ::tex::gamma_mode; +import ::state::normal; +import ::base::mono_mode; + +export material Plastic_Rough_Black_A(*) + = OmniPBR( + diffuse_color_constant: color(0.2, 0.2, 0.2), + diffuse_texture: texture_2d("../Textures/T_Plastic_Rough_Black_A_Albedo.png", ::tex::gamma_srgb), + albedo_desaturation: 0.0, + albedo_add: 0.0, + albedo_brightness: 1.0, + diffuse_tint: color(1, 1, 1), + reflection_roughness_constant: 0.5, + reflection_roughness_texture_influence: 1.0, + reflectionroughness_texture: texture_2d(), + metallic_constant: 0.0, + metallic_texture_influence: 1.0, + metallic_texture: texture_2d(), + specular_level: 0.5, + enable_ORM_texture: true, + ORM_texture: texture_2d("../Textures/T_Plastic_Rough_Black_A_ORM.png", ::tex::gamma_linear), + ao_to_diffuse: 0.0, + ao_texture: texture_2d(), + enable_emission: false, + emissive_color: color(1, 0.1, 0.1), + emissive_color_texture: texture_2d(), + emissive_mask_texture: texture_2d(), + emissive_intensity: 40.0, + enable_opacity: false, + enable_opacity_texture: false, + opacity_constant: 1.0, + opacity_texture: texture_2d(), + opacity_mode: ::base::mono_average, + opacity_threshold: 0.0, + bump_factor: 2.5, + normalmap_texture: texture_2d("../Textures/T_Plastic_Rough_Black_A_Normal.png", ::tex::gamma_linear), + detail_bump_factor: 0.30000001192092896, + detail_normalmap_texture: texture_2d(), + flip_tangent_u: false, + flip_tangent_v: true, + project_uvw: false, + world_or_object: false, + uv_space_index: 0, + texture_translate: float2(0, 0), + texture_rotate: 0.0, + texture_scale: float2(1, 1), + detail_texture_translate: float2(0, 0), + detail_texture_rotate: 0.0, + detail_texture_scale: float2(1, 1)); diff --git a/assets/conveyors/Material Library/Plastic/Textures/T_Acrylic_Clear_Glossy_A_Rough.png b/assets/conveyors/Material Library/Plastic/Textures/T_Acrylic_Clear_Glossy_A_Rough.png new file mode 100644 index 0000000..6bd628b Binary files /dev/null and b/assets/conveyors/Material Library/Plastic/Textures/T_Acrylic_Clear_Glossy_A_Rough.png differ diff --git a/assets/conveyors/Material Library/Plastic/Textures/T_Plastic_Rough_Black_A_Normal.png b/assets/conveyors/Material Library/Plastic/Textures/T_Plastic_Rough_Black_A_Normal.png new file mode 100644 index 0000000..b85de8a Binary files /dev/null and b/assets/conveyors/Material Library/Plastic/Textures/T_Plastic_Rough_Black_A_Normal.png differ diff --git a/assets/conveyors/Material%20Library/physics_material.usd b/assets/conveyors/Material%20Library/physics_material.usd new file mode 100644 index 0000000..f805520 Binary files /dev/null and b/assets/conveyors/Material%20Library/physics_material.usd differ diff --git a/assets/conveyors/Textures/Acrylic_Clear_Glossy_A.usd b/assets/conveyors/Textures/Acrylic_Clear_Glossy_A.usd new file mode 100644 index 0000000..8693f58 Binary files /dev/null and b/assets/conveyors/Textures/Acrylic_Clear_Glossy_A.usd differ diff --git a/assets/conveyors/Textures/M_ConveyorBelt_A01_Belt.usd b/assets/conveyors/Textures/M_ConveyorBelt_A01_Belt.usd new file mode 100644 index 0000000..634bffd Binary files /dev/null and b/assets/conveyors/Textures/M_ConveyorBelt_A01_Belt.usd differ diff --git a/assets/conveyors/Textures/M_ConveyorBelt_A01_Decal.usd b/assets/conveyors/Textures/M_ConveyorBelt_A01_Decal.usd new file mode 100644 index 0000000..bb62b59 Binary files /dev/null and b/assets/conveyors/Textures/M_ConveyorBelt_A01_Decal.usd differ diff --git a/assets/conveyors/Textures/MetalPainted_Black_Glossy_A.usd b/assets/conveyors/Textures/MetalPainted_Black_Glossy_A.usd new file mode 100644 index 0000000..5ee82b8 Binary files /dev/null and b/assets/conveyors/Textures/MetalPainted_Black_Glossy_A.usd differ diff --git a/assets/conveyors/Textures/MetalPainted_Blue_Glossy_A.usd b/assets/conveyors/Textures/MetalPainted_Blue_Glossy_A.usd new file mode 100644 index 0000000..6c1bb9a Binary files /dev/null and b/assets/conveyors/Textures/MetalPainted_Blue_Glossy_A.usd differ diff --git a/assets/conveyors/Textures/Metal_Rough_A.usd b/assets/conveyors/Textures/Metal_Rough_A.usd new file mode 100644 index 0000000..d020df9 Binary files /dev/null and b/assets/conveyors/Textures/Metal_Rough_A.usd differ diff --git a/assets/conveyors/Textures/Plastic_Orange_A.usd b/assets/conveyors/Textures/Plastic_Orange_A.usd new file mode 100644 index 0000000..55383a1 Binary files /dev/null and b/assets/conveyors/Textures/Plastic_Orange_A.usd differ diff --git a/assets/conveyors/Textures/Plastic_Red_A.usd b/assets/conveyors/Textures/Plastic_Red_A.usd new file mode 100644 index 0000000..e39c18d Binary files /dev/null and b/assets/conveyors/Textures/Plastic_Red_A.usd differ diff --git a/assets/conveyors/Textures/Plastic_Rough_Black_A.usd b/assets/conveyors/Textures/Plastic_Rough_Black_A.usd new file mode 100644 index 0000000..1195640 Binary files /dev/null and b/assets/conveyors/Textures/Plastic_Rough_Black_A.usd differ diff --git a/assets/conveyors/Textures/Steel_A.usd b/assets/conveyors/Textures/Steel_A.usd new file mode 100644 index 0000000..294600c Binary files /dev/null and b/assets/conveyors/Textures/Steel_A.usd differ diff --git a/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_Albedo.png b/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_Albedo.png new file mode 100644 index 0000000..6a1407c Binary files /dev/null and b/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_Albedo.png differ diff --git a/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_Normal.png b/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_Normal.png new file mode 100644 index 0000000..f935d22 Binary files /dev/null and b/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_Normal.png differ diff --git a/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_ORM.png b/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_ORM.png new file mode 100644 index 0000000..de82c8c Binary files /dev/null and b/assets/conveyors/Textures/T_ConveyorBelt_A01_Belt_ORM.png differ diff --git a/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_Albedo.png b/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_Albedo.png new file mode 100644 index 0000000..3c27b3e Binary files /dev/null and b/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_Albedo.png differ diff --git a/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_Alpha.png b/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_Alpha.png new file mode 100644 index 0000000..4a59717 Binary files /dev/null and b/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_Alpha.png differ diff --git a/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_ORM.png b/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_ORM.png new file mode 100644 index 0000000..4183089 Binary files /dev/null and b/assets/conveyors/Textures/T_ConveyorsBelt_A01_Decal_ORM.png differ diff --git a/cv/circular_section.py b/cv/circular_section.py new file mode 100644 index 0000000..9c9389f --- /dev/null +++ b/cv/circular_section.py @@ -0,0 +1,99 @@ +""" +Circular-cross-section criterion (spec Category D "не подходит для сортировки без доупаковки"). + +K = r_inscribed / R_circumscribed for a cross-section outline; an object "has a circle in +section" when K > 0.8 in ANY of its principal cross-sections. Implemented so it works on a +(partial) point cloud measured under the dimension-estimation cameras: the circle CENTRE is +fit to the boundary (Kasa) so an occluded-bottom arc still reads correctly, and a low +circle-fit residual is required so an elongated rounded blob is not mistaken for a circle. + +This module both (a) exposes `circular_section_K(points)` for the pipeline and (b) validates +it against ground truth on every mesh, including the ones that are NOT round (boxes, etc.). +Run inside Isaac Sim via isaacsim_send.py (reads /World/CVObjects mesh geometry). +""" +import omni.usd, numpy as np, json +from pxr import UsdGeom, Usd, Gf +stage=omni.usd.get_context().get_stage() +K_ROUND=0.80 # spec threshold + +def _kasa(P): + x,y=P[:,0],P[:,1]; A=np.c_[2*x,2*y,np.ones(len(x))]; b=x*x+y*y + s,*_=np.linalg.lstsq(A,b,rcond=None); cx,cy,cc=s; r=np.sqrt(max(cc+cx*cx+cy*cy,1e-12)) + return cx,cy,r,np.abs(np.hypot(x-cx,y-cy)-r).mean() + +def section_K(xy): + """True r_inscribed/R_circumscribed of a full cross-section outline via the convex-hull + incenter (largest inscribed circle) and the circumscribed radius from that centre. + K=1 for a circle, b/a for an ellipse, 0.707 for a square, short/long for a rectangle.""" + from scipy.spatial import ConvexHull + if len(xy)<20: return None,0.0,1.0 + try: h=ConvexHull(xy) + except Exception: return None,0.0,1.0 + V=xy[h.vertices] # CCW hull vertices + A=V; B=np.roll(V,-1,axis=0); E=B-A; L=np.linalg.norm(E,axis=1)+1e-12 + mn=xy.min(0); mx=xy.max(0) + G=np.stack(np.meshgrid(np.linspace(mn[0],mx[0],40),np.linspace(mn[1],mx[1],40)),-1).reshape(-1,2) + # signed distance from each grid point to each hull edge (CCW -> interior side positive) + d=(E[:,0][None,:]*(G[:,1][:,None]-A[:,1][None,:]) - E[:,1][None,:]*(G[:,0][:,None]-A[:,0][None,:]))/L[None,:] + inside=(d>0).all(1) + if inside.sum()<3: return 0.0,1.0,1.0 + rin=float(d[inside].min(1).max()) # max inscribed circle radius (its own centre) + # min enclosing circle radius (its own centre): grid centre minimising max distance to hull vertices + Gd=np.stack(np.meshgrid(np.linspace(mn[0],mx[0],48),np.linspace(mn[1],mx[1],48)),-1).reshape(-1,2) + Rout=float(np.linalg.norm(Gd[:,None,:]-V[None,:,:],axis=2).max(1).min()) + return rin/max(Rout,1e-9),1.0,0.0 + +def circular_section_K(points): + """Max K over cross-sections sampled along each principal axis (a circle in ANY section + -> round). Returns (max_K, is_round, best_section).""" + if len(points)<60: return 0.0,False,None + c=points.mean(0); Q=points-c; _,_,V=np.linalg.svd(Q,full_matrices=False); proj=Q@V.T + best=0.0; best_sec=None + for a in range(3): + o=[i for i in range(3) if i!=a]; ca=proj[:,a]; sp=np.ptp(ca)+1e-9 + for frac in (0.25,0.375,0.5,0.625,0.75): # sample slices along the axis + lvl=np.percentile(ca,frac*100) + sl=proj[np.abs(ca-lvl)<0.07*sp][:,o] + K,cov,rr=section_K(sl) + if K is not None and rr<0.15 and K>best: + best=K; best_sec=(a,round(frac,2),round(cov,2),rr) + return round(best,3), (best>K_ROUND), best_sec + +# ---------- validate on all meshes (full GT geometry) ---------- +def mesh_points(nm, nsamp=40000): + """Dense, uniform surface sample (barycentric) so thin cross-sections are well populated.""" + root=stage.GetPrimAtPath(f"/World/CVObjects/{nm}"); rng=np.random.default_rng(0); out=[] + for m in Usd.PrimRange(root): + if m.GetTypeName()!="Mesh": continue + P=np.array(UsdGeom.Mesh(m).GetPointsAttr().Get(),dtype=np.float64) + idx=np.array(UsdGeom.Mesh(m).GetFaceVertexIndicesAttr().Get()) + if len(idx)%3: continue + tris=P[idx].reshape(-1,3,3); v0,v1,v2=tris[:,0],tris[:,1],tris[:,2] + area=0.5*np.linalg.norm(np.cross(v1-v0,v2-v0),axis=1); s=area.sum() + if s<=0: continue + ti=rng.choice(len(tris),nsamp,p=area/s) + r1=np.sqrt(rng.random(nsamp)); r2=rng.random(nsamp) + out.append(((1-r1)[:,None]*v0[ti]+(r1*(1-r2))[:,None]*v1[ti]+(r1*r2)[:,None]*v2[ti])) + return np.concatenate(out) if out else np.zeros((0,3)) + +objs=[p.GetName() for p in stage.GetPrimAtPath("/World/CVObjects").GetChildren()] +rows=[]; tp=fp=tn=fn=0 +for nm in sorted(objs): + mp=stage.GetPrimAtPath(f"/World/CVObjects/{nm}/Mesh") + gt_k=mp.GetCustomDataByKey("gt_k_round"); gt_zone=mp.GetCustomDataByKey("gt_zone") + P=mesh_points(nm) + maxK,is_round,sec=circular_section_K(P) + gt_round=(gt_k is not None and gt_k>0.80) + ok = (is_round==gt_round) + if is_round and gt_round: tp+=1 + elif is_round and not gt_round: fp+=1 + elif not is_round and not gt_round: tn+=1 + else: fn+=1 + rows.append((nm,gt_zone,round(gt_k,3) if gt_k else None,maxK,is_round,gt_round,ok)) + print(f" {nm:16s} zone={gt_zone} GTk={gt_k:.2f} measuredK={maxK:.2f} round={str(is_round):5s} GTround={str(gt_round):5s} {'OK' if ok else 'MISS'}") +prec=tp/(tp+fp) if tp+fp else 0; rec=tp/(tp+fn) if tp+fn else 0 +print(f"\ncircle-in-section detection vs GT k_round>0.8:") +print(f" round: precision={prec:.2f} recall={rec:.2f} | TP={tp} FP={fp} TN={tn} FN={fn} acc={(tp+tn)/len(rows):.2f} ({tp+tn}/{len(rows)})") +non_round=[r for r in rows if not r[5]] +print(f" non-round objects correctly rejected: {sum(1 for r in non_round if not r[4])}/{len(non_round)} ({[r[0] for r in non_round if r[4]]} wrongly flagged)") +json.dump([list(r) for r in rows],open("/home/dasha/isaac_assets/cv/circular_section_results.json","w"),indent=2) diff --git a/cv/circular_section_results.json b/cv/circular_section_results.json new file mode 100644 index 0000000..159c0fe --- /dev/null +++ b/cv/circular_section_results.json @@ -0,0 +1,281 @@ +[ + [ + "backpack", + "C", + 0.82, + 0.795, + false, + true, + false + ], + [ + "bag", + "D", + 0.896, + 0.889, + true, + true, + true + ], + [ + "banana", + "D", + 0.94, + 0.917, + true, + true, + true + ], + [ + "barrel", + "C", + 0.995, + 0.933, + true, + true, + true + ], + [ + "bolts_cluster", + "B", + 0.718, + 0.708, + false, + false, + true + ], + [ + "bottle", + "D", + 0.995, + 0.935, + true, + true, + true + ], + [ + "box_300x200x200", + "B", + 0.72, + 0.685, + false, + false, + true + ], + [ + "box_400x400x300", + "C", + 0.716, + 0.688, + false, + false, + true + ], + [ + "broom", + "C", + 0.966, + 0.552, + false, + true, + false + ], + [ + "bucket", + "D", + 0.995, + 0.934, + true, + true, + true + ], + [ + "chip_bag", + "D", + 0.811, + 0.589, + false, + true, + false + ], + [ + "cone", + "C", + 0.991, + 0.931, + true, + true, + true + ], + [ + "cylinder", + "D", + 0.867, + 0.705, + false, + true, + false + ], + [ + "detergent", + "B", + 0.742, + 0.684, + false, + false, + true + ], + [ + "headphones", + "D", + 0.807, + 0.76, + false, + true, + false + ], + [ + "helmet", + "D", + 0.895, + 0.803, + true, + true, + true + ], + [ + "lunchbox", + "B", + 0.646, + 0.619, + false, + false, + true + ], + [ + "mug", + "D", + 0.985, + 0.858, + true, + true, + true + ], + [ + "office_chair", + "C", + 0.996, + 0.791, + false, + true, + false + ], + [ + "pallet", + "C", + 0.711, + 0.685, + false, + false, + true + ], + [ + "parcel_box", + "B", + 0.699, + 0.675, + false, + false, + true + ], + [ + "pen", + "C", + 0.842, + 0.828, + true, + true, + true + ], + [ + "perfume", + "D", + 0.924, + 0.878, + true, + true, + true + ], + [ + "pillow", + "C", + 0.905, + 0.698, + false, + true, + false + ], + [ + "plate", + "D", + 0.998, + 0.936, + true, + true, + true + ], + [ + "pouf", + "C", + 0.994, + 0.935, + true, + true, + true + ], + [ + "sneaker", + "B", + 0.706, + 0.69, + false, + false, + true + ], + [ + "tire", + "C", + 0.994, + 0.933, + true, + true, + true + ], + [ + "tool_case", + "B", + 0.454, + 0.43, + false, + false, + true + ], + [ + "umbrella", + "C", + 0.628, + 0.879, + true, + false, + false + ], + [ + "watch", + "C", + 0.995, + 0.668, + false, + true, + false + ] +] \ No newline at end of file diff --git a/cv/crestereo/.gitattributes b/cv/crestereo/.gitattributes new file mode 100644 index 0000000..dfe0770 --- /dev/null +++ b/cv/crestereo/.gitattributes @@ -0,0 +1,2 @@ +# Auto detect text files and perform LF normalization +* text=auto diff --git a/cv/crestereo/.gitignore b/cv/crestereo/.gitignore new file mode 100644 index 0000000..0514910 --- /dev/null +++ b/cv/crestereo/.gitignore @@ -0,0 +1,156 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintainted in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +vis_results/ +models/* +test_data/* \ No newline at end of file diff --git a/cv/crestereo/README.md b/cv/crestereo/README.md new file mode 100644 index 0000000..4803b23 --- /dev/null +++ b/cv/crestereo/README.md @@ -0,0 +1,46 @@ +# CREStereo-Pytorch + Non-official Pytorch implementation of the CREStereo (CVPR 2022 Oral) model converted from the original MegEngine implementation. + +![!CREStereo-Pytorch stereo detph estimation](https://github.com/ibaiGorordo/CREStereo-Pytorch/blob/main/doc/img/output.jpg) + +**update 2023/01/03**: +- enable DistributedDataParallel (DDP) training, training time is much faster than before. + +```shell +# train DDP +# change 'dist' to True in /cfgs/train.yaml file +python -m torch.distributed.launch --nproc_per_node=8 train.py +# train DP +# change 'dist' to False in /cfgs/train.yaml file +python train.py +``` + +# Important +- This is just an effort to try to implement the CREStereo model into Pytorch from MegEngine due to the issues of the framework to convert to other formats (https://github.com/megvii-research/CREStereo/issues/3). +- I am not the author of the paper, and I am don't fully understand what the model is doing. Therefore, there might be small differences with the original model that might impact the performance. +- I have not added any license, since the repository uses code from different repositories. Check the License section below for more detail. + +# Pretrained model +- Download the model from [here](https://drive.google.com/file/d/1D2s1v4VhJlNz98FQpFxf_kBAKQVN_7xo/view?usp=sharing) and save it into the **[models](https://github.com/ibaiGorordo/CREStereo-Pytorch/tree/main/models)** folder. +- The model was converted from the original **[MegEngine weights](https://drive.google.com/file/d/1Wx_-zDQh7BUFBmN9im_26DFpnf3AkXj4/view)** using the `convert_weights.py` script. Place the MegEngine weights (crestereo_eth3d.mge) file into the **[models](https://github.com/ibaiGorordo/CREStereo-Pytorch/tree/main/models)** folder before the conversion. + +# ONNX Conversion +- After either downloading the pretrained weights or training your own model, you will have a `models/crestereo_eth3d.pth` file. If you want to run your model with ONNX, you need to run the convert_to_onnx.py script. The script has two parts: +1. Convert the model to an ONNX model that takes in left, right images as well as an initial flow estimate (takes a few seconds) +2. Convert the model to an ONNX model that takes in left, right images and NO initial flow estimate (takes several minutes and requires pytorch >= 1.12) +(afaik) You will need both models to get the same results as you do from test_model.py. +- Run the test_onnx_model.py script to verify your models work as expected! +- NOTE: although the test_model.py script works with any size images as input, once you have converted your +Pytorch model into ONNX models, you must provide them with the image sizes used at conversion time or it will not work. + +# Licences: +- CREStereo (Apache License 2.0): https://github.com/megvii-research/CREStereo/blob/master/LICENSE +- RAFT (BSD 3-Clause):https://github.com/princeton-vl/RAFT/blob/master/LICENSE +- LoFTR (Apache License 2.0):https://github.com/zju3dv/LoFTR/blob/master/LICENSE + +# References: +- CREStereo: https://github.com/megvii-research/CREStereo +- RAFT: https://github.com/princeton-vl/RAFT +- LoFTR: https://github.com/zju3dv/LoFTR +- Grid sample replacement: https://zenn.dev/pinto0309/scraps/7d4032067d0160 +- torch2mge: https://github.com/MegEngine/torch2mge diff --git a/cv/crestereo/cfgs/train.yaml b/cv/crestereo/cfgs/train.yaml new file mode 100644 index 0000000..6f1e5d7 --- /dev/null +++ b/cv/crestereo/cfgs/train.yaml @@ -0,0 +1,20 @@ +seed: 0 +mixed_precision: false +base_lr: 4.0e-4 + +nr_gpus: 8 +batch_size: 4 +n_total_epoch: 600 +minibatch_per_epoch: 500 + +loadmodel: ~ +log_dir: "./train_log" +model_save_freq_epoch: 1 + +max_disp: 256 +image_width: 512 +image_height: 384 +training_data_path: "./stereo_trainset/crestereo" + +log_level: "logging.INFO" +dist: True # True for DDP, False for DP diff --git a/cv/crestereo/convert_to_onnx.py b/cv/crestereo/convert_to_onnx.py new file mode 100644 index 0000000..cf20b82 --- /dev/null +++ b/cv/crestereo/convert_to_onnx.py @@ -0,0 +1,49 @@ +import torch +import torch.nn.functional as F +import numpy as np +import cv2 +from imread_from_url import imread_from_url + +from nets import Model + +if __name__ == '__main__': + + model_path = "models/crestereo_eth3d.pth" + + model = Model(max_disp=256, mixed_precision=False, test_mode=True) + model.load_state_dict(torch.load(model_path), strict=True) + model.eval() + + in_h, in_w = (480, 640) + t1_half = torch.rand(1, 3, in_h//2, in_w//2) + t2_half = torch.rand(1, 3, in_h//2, in_w//2) + + t1 = torch.rand(1, 3, in_h, in_w) + t2 = torch.rand(1, 3, in_h, in_w) + flow_init = torch.rand(1, 2, in_h//2, in_w//2) + + # Export the model + torch.onnx.export(model, + (t1, t2, flow_init), + "crestereo.onnx", # where to save the model (can be a file or file-like object) + export_params=True, # store the trained parameter weights inside the model file + opset_version=12, # the ONNX version to export the model to + do_constant_folding=True, # whether to execute constant folding for optimization + input_names = ['left', 'right','flow_init'], # the model's input names + output_names = ['output']) + + # Export the model without init_flow (it takes a lot of time) + # !! Does not work prior to pytorch 1.12 (confirmed working on pytorch 2.0.0) + # Ref: https://github.com/pytorch/pytorch/pull/73760 + torch.onnx.export(model, + (t1_half, t2_half), + "crestereo_without_flow.onnx", # where to save the model (can be a file or file-like object) + export_params=True, # store the trained parameter weights inside the model file + opset_version=12, # the ONNX version to export the model to + do_constant_folding=True, # whether to execute constant folding for optimization + input_names = ['left', 'right'], # the model's input names + output_names = ['output']) + + + + diff --git a/cv/crestereo/convert_weights.py b/cv/crestereo/convert_weights.py new file mode 100644 index 0000000..c765f39 --- /dev/null +++ b/cv/crestereo/convert_weights.py @@ -0,0 +1,26 @@ +import copy +import torch +import numpy as np +import megengine as mge + +from nets import Model + +# Read Megengine parameters +pretrained_dict = mge.load("models/crestereo_eth3d.mge") + +model = Model(max_disp=256, mixed_precision=False, test_mode=True) +model.eval() + +state_dict = model.state_dict() +for key, value in pretrained_dict['state_dict'].items(): + + print(f"Converting {key}") + # Fix shape mismatch + if value.shape[0] == 1: + value = np.squeeze(value) + + state_dict[key] = torch.tensor(value) + +output_path = "models/crestereo_eth3d.pth" +torch.save(state_dict, output_path) +print(f"\nModel saved to: {output_path}") \ No newline at end of file diff --git a/cv/crestereo/dataset.py b/cv/crestereo/dataset.py new file mode 100644 index 0000000..67a62e1 --- /dev/null +++ b/cv/crestereo/dataset.py @@ -0,0 +1,215 @@ +import os +import cv2 +import glob +import numpy as np +from PIL import Image, ImageEnhance + +from torch.utils.data import Dataset + + +class Augmentor: + def __init__( + self, + image_height=384, + image_width=512, + max_disp=256, + scale_min=0.6, + scale_max=1.0, + seed=0, + ): + super().__init__() + self.image_height = image_height + self.image_width = image_width + self.max_disp = max_disp + self.scale_min = scale_min + self.scale_max = scale_max + self.rng = np.random.RandomState(seed) + + def chromatic_augmentation(self, img): + random_brightness = np.random.uniform(0.8, 1.2) + random_contrast = np.random.uniform(0.8, 1.2) + random_gamma = np.random.uniform(0.8, 1.2) + + img = Image.fromarray(img) + + enhancer = ImageEnhance.Brightness(img) + img = enhancer.enhance(random_brightness) + enhancer = ImageEnhance.Contrast(img) + img = enhancer.enhance(random_contrast) + + gamma_map = [ + 255 * 1.0 * pow(ele / 255.0, random_gamma) for ele in range(256) + ] * 3 + img = img.point(gamma_map) # use PIL's point-function to accelerate this part + + img_ = np.array(img) + + return img_ + + def __call__(self, left_img, right_img, left_disp): + # 1. chromatic augmentation + left_img = self.chromatic_augmentation(left_img) + right_img = self.chromatic_augmentation(right_img) + + # 2. spatial augmentation + # 2.1) rotate & vertical shift for right image + if self.rng.binomial(1, 0.5): + angle, pixel = 0.1, 2 + px = self.rng.uniform(-pixel, pixel) + ag = self.rng.uniform(-angle, angle) + image_center = ( + self.rng.uniform(0, right_img.shape[0]), + self.rng.uniform(0, right_img.shape[1]), + ) + rot_mat = cv2.getRotationMatrix2D(image_center, ag, 1.0) + right_img = cv2.warpAffine( + right_img, rot_mat, right_img.shape[1::-1], flags=cv2.INTER_LINEAR + ) + trans_mat = np.float32([[1, 0, 0], [0, 1, px]]) + right_img = cv2.warpAffine( + right_img, trans_mat, right_img.shape[1::-1], flags=cv2.INTER_LINEAR + ) + + # 2.2) random resize + resize_scale = self.rng.uniform(self.scale_min, self.scale_max) + + left_img = cv2.resize( + left_img, + None, + fx=resize_scale, + fy=resize_scale, + interpolation=cv2.INTER_LINEAR, + ) + right_img = cv2.resize( + right_img, + None, + fx=resize_scale, + fy=resize_scale, + interpolation=cv2.INTER_LINEAR, + ) + + disp_mask = (left_disp < float(self.max_disp / resize_scale)) & (left_disp > 0) + disp_mask = disp_mask.astype("float32") + disp_mask = cv2.resize( + disp_mask, + None, + fx=resize_scale, + fy=resize_scale, + interpolation=cv2.INTER_LINEAR, + ) + + left_disp = ( + cv2.resize( + left_disp, + None, + fx=resize_scale, + fy=resize_scale, + interpolation=cv2.INTER_LINEAR, + ) + * resize_scale + ) + + # 2.3) random crop + h, w, c = left_img.shape + dx = w - self.image_width + dy = h - self.image_height + dy = self.rng.randint(min(0, dy), max(0, dy) + 1) + dx = self.rng.randint(min(0, dx), max(0, dx) + 1) + + M = np.float32([[1.0, 0.0, -dx], [0.0, 1.0, -dy]]) + left_img = cv2.warpAffine( + left_img, + M, + (self.image_width, self.image_height), + flags=cv2.INTER_LINEAR, + borderValue=0, + ) + right_img = cv2.warpAffine( + right_img, + M, + (self.image_width, self.image_height), + flags=cv2.INTER_LINEAR, + borderValue=0, + ) + left_disp = cv2.warpAffine( + left_disp, + M, + (self.image_width, self.image_height), + flags=cv2.INTER_LINEAR, + borderValue=0, + ) + disp_mask = cv2.warpAffine( + disp_mask, + M, + (self.image_width, self.image_height), + flags=cv2.INTER_LINEAR, + borderValue=0, + ) + + # 3. add random occlusion to right image + if self.rng.binomial(1, 0.5): + sx = int(self.rng.uniform(50, 100)) + sy = int(self.rng.uniform(50, 100)) + cx = int(self.rng.uniform(sx, right_img.shape[0] - sx)) + cy = int(self.rng.uniform(sy, right_img.shape[1] - sy)) + right_img[cx - sx : cx + sx, cy - sy : cy + sy] = np.mean( + np.mean(right_img, 0), 0 + )[np.newaxis, np.newaxis] + + return left_img, right_img, left_disp, disp_mask + + +class CREStereoDataset(Dataset): + def __init__(self, root): + super().__init__() + self.imgs = glob.glob(os.path.join(root, "**/*_left.jpg"), recursive=True) + self.augmentor = Augmentor( + image_height=384, + image_width=512, + max_disp=256, + scale_min=0.6, + scale_max=1.0, + seed=0, + ) + self.rng = np.random.RandomState(0) + + def get_disp(self, path): + disp = cv2.imread(path, cv2.IMREAD_UNCHANGED) + return disp.astype(np.float32) / 32 + + def __getitem__(self, index): + # find path + left_path = self.imgs[index] + prefix = left_path[: left_path.rfind("_")] + right_path = prefix + "_right.jpg" + left_disp_path = prefix + "_left.disp.png" + right_disp_path = prefix + "_right.disp.png" + + # read img, disp + left_img = cv2.imread(left_path, cv2.IMREAD_COLOR) + right_img = cv2.imread(right_path, cv2.IMREAD_COLOR) + left_disp = self.get_disp(left_disp_path) + right_disp = self.get_disp(right_disp_path) + + if self.rng.binomial(1, 0.5): + left_img, right_img = np.fliplr(right_img), np.fliplr(left_img) + left_disp, right_disp = np.fliplr(right_disp), np.fliplr(left_disp) + left_disp[left_disp == np.inf] = 0 + + # augmentaion + left_img, right_img, left_disp, disp_mask = self.augmentor( + left_img, right_img, left_disp + ) + + left_img = left_img.transpose(2, 0, 1).astype("uint8") + right_img = right_img.transpose(2, 0, 1).astype("uint8") + + return { + "left": left_img, + "right": right_img, + "disparity": left_disp, + "mask": disp_mask, + } + + def __len__(self): + return len(self.imgs) diff --git a/cv/crestereo/doc/img/output.jpg b/cv/crestereo/doc/img/output.jpg new file mode 100644 index 0000000..c7dc0ec Binary files /dev/null and b/cv/crestereo/doc/img/output.jpg differ diff --git a/cv/crestereo/function_convertion_tests/test_bilinear_sampler.py b/cv/crestereo/function_convertion_tests/test_bilinear_sampler.py new file mode 100644 index 0000000..2e12267 --- /dev/null +++ b/cv/crestereo/function_convertion_tests/test_bilinear_sampler.py @@ -0,0 +1,44 @@ +import pickle +import numpy as np +import megengine as mge + +import torch +import torch.nn.functional as F + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + xgrid, ygrid = coords.split([1,1], dim=-1) + xgrid = 2*xgrid/(W-1) - 1 + ygrid = 2*ygrid/(H-1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + img = F.grid_sample(img, grid, align_corners=True) + + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + +def test_bilinear_sampler(): + # Getting back the megengine objects: + with open('test_data/bilinear_sampler_test.pickle', 'rb') as f: + right_feature_prev, coords, right_feature = pickle.load(f) + + right_feature_prev = torch.tensor(right_feature_prev.numpy()) + coords = torch.tensor(coords.numpy()) + right_feature = right_feature.numpy() + + # Test Pytorch + right_feature_pytorch = bilinear_sampler(right_feature_prev, coords).numpy() + + error = np.mean(right_feature_pytorch-right_feature) + print(f"test_coords_grid - Avg. Error: {error}, \n \ + Original shape: {coords.numpy().shape},\n \ + Obtained shape: {right_feature_pytorch.shape}, Expected shape: {right_feature.shape}") + +if __name__ == '__main__': + + test_bilinear_sampler() \ No newline at end of file diff --git a/cv/crestereo/function_convertion_tests/test_coords_grid.py b/cv/crestereo/function_convertion_tests/test_coords_grid.py new file mode 100644 index 0000000..0f597f4 --- /dev/null +++ b/cv/crestereo/function_convertion_tests/test_coords_grid.py @@ -0,0 +1,29 @@ +import pickle +import numpy as np +import megengine as mge + +import torch +import torch.nn.functional as F + +def coords_grid(batch, ht, wd, device): + coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device), indexing='ij') + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + +def test_coords_grid(): + # Getting back the megengine objects: + with open('test_data/coords_grid_test.pickle', 'rb') as f: + batch, ht, wd, coords = pickle.load(f) + + coords = coords.numpy() + + # Test Pytorch + coords_pytorch = coords_grid(batch, ht, wd, 'cpu').numpy() + + error = np.mean(coords_pytorch-coords) + print(f"test_coords_grid - Avg. Error: {error}, \n \ + Obtained shape: {coords_pytorch.shape}, Expected shape: {coords.shape}") + +if __name__ == '__main__': + + test_coords_grid() \ No newline at end of file diff --git a/cv/crestereo/function_convertion_tests/test_data/bilinear_sampler_test.pickle b/cv/crestereo/function_convertion_tests/test_data/bilinear_sampler_test.pickle new file mode 100644 index 0000000..7156e08 Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/bilinear_sampler_test.pickle differ diff --git a/cv/crestereo/function_convertion_tests/test_data/coords_grid_test.pickle b/cv/crestereo/function_convertion_tests/test_data/coords_grid_test.pickle new file mode 100644 index 0000000..18b3799 Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/coords_grid_test.pickle differ diff --git a/cv/crestereo/function_convertion_tests/test_data/manual_pad_test0_4.pickle b/cv/crestereo/function_convertion_tests/test_data/manual_pad_test0_4.pickle new file mode 100644 index 0000000..37dddcb Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/manual_pad_test0_4.pickle differ diff --git a/cv/crestereo/function_convertion_tests/test_data/manual_pad_test1_1.pickle b/cv/crestereo/function_convertion_tests/test_data/manual_pad_test1_1.pickle new file mode 100644 index 0000000..8fb52a8 Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/manual_pad_test1_1.pickle differ diff --git a/cv/crestereo/function_convertion_tests/test_data/meshgrid_np_test.pkl b/cv/crestereo/function_convertion_tests/test_data/meshgrid_np_test.pkl new file mode 100644 index 0000000..c3f9cb3 Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/meshgrid_np_test.pkl differ diff --git a/cv/crestereo/function_convertion_tests/test_data/offset_test.pkl b/cv/crestereo/function_convertion_tests/test_data/offset_test.pkl new file mode 100644 index 0000000..b92d22e Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/offset_test.pkl differ diff --git a/cv/crestereo/function_convertion_tests/test_data/split_test.pkl b/cv/crestereo/function_convertion_tests/test_data/split_test.pkl new file mode 100644 index 0000000..e5c96a1 Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/split_test.pkl differ diff --git a/cv/crestereo/function_convertion_tests/test_data/split_test_list.pkl b/cv/crestereo/function_convertion_tests/test_data/split_test_list.pkl new file mode 100644 index 0000000..68692b2 Binary files /dev/null and b/cv/crestereo/function_convertion_tests/test_data/split_test_list.pkl differ diff --git a/cv/crestereo/function_convertion_tests/test_manual_pad.py b/cv/crestereo/function_convertion_tests/test_manual_pad.py new file mode 100644 index 0000000..c5a1bbe --- /dev/null +++ b/cv/crestereo/function_convertion_tests/test_manual_pad.py @@ -0,0 +1,51 @@ +import pickle +import numpy as np +import megengine as mge + +import torch +import torch.nn.functional as F + +def manual_pad(x, pady, padx): + + pad = (padx, padx, pady, pady) + return F.pad(torch.tensor(x), pad, "replicate") + + +def test_pad_1_1(): + # Getting back the megengine objects: + with open('test_data/manual_pad_test1_1.pickle', 'rb') as f: + right_feature, pady, padx, right_pad = pickle.load(f) + + right_feature = right_feature.numpy() + right_pad = right_pad.numpy() + + # Test Pytorch + right_pad_pytorch = manual_pad(right_feature, pady, padx).numpy() + + error = np.mean(right_pad_pytorch-right_pad) + print(f"test_pad_1_1 - Avg. Error: {error}, \n \ + Orig. shape: {right_feature.shape}, \n \ + Padded shape: {right_pad_pytorch.shape}, Expected shape: {right_pad.shape}") + +def test_pad_0_4(): + # Getting back the megengine objects: + with open('test_data/manual_pad_test0_4.pickle', 'rb') as f: + right_feature, pady, padx, right_pad = pickle.load(f) + + right_feature = right_feature.numpy() + right_pad = right_pad.numpy() + + # Test Pytorch + right_pad_pytorch = manual_pad(right_feature, pady, padx).numpy() + + error = np.mean(right_pad_pytorch-right_pad) + print(f"test_pad_0_4 - Avg. Error: {error}, \n \ + Orig. shape: {right_feature.shape}, \n \ + Padded shape: {right_pad_pytorch.shape}, Expected shape: {right_pad.shape}") + + +if __name__ == '__main__': + + test_pad_1_1() + + test_pad_0_4() \ No newline at end of file diff --git a/cv/crestereo/function_convertion_tests/test_meshgrid.py b/cv/crestereo/function_convertion_tests/test_meshgrid.py new file mode 100644 index 0000000..9883641 --- /dev/null +++ b/cv/crestereo/function_convertion_tests/test_meshgrid.py @@ -0,0 +1,30 @@ +import pickle +import numpy as np +import megengine as mge + +import torch +import torch.nn.functional as F + +def test_meshgrid(): + # Getting back the megengine objects: + with open('test_data/meshgrid_np_test.pkl', 'rb') as f: + rx, dilatex, ry, dilatey, x_grid, y_grid = pickle.load(f) + + x_grid = x_grid.numpy() + y_grid = y_grid.numpy() + + # Test Pytorch + x_grid_pytorch, y_grid_pytorch = torch.meshgrid(torch.arange(-rx, rx + 1, dilatex, device='cpu'), + torch.arange(-ry, ry + 1, dilatey, device='cpu'), indexing='xy') + + + error_x = np.mean(x_grid_pytorch.numpy()-x_grid) + error_y = np.mean(y_grid_pytorch.numpy()-y_grid) + print(f"test_meshgrid (X) - Avg. Error: {error_x}, \n \ + Obtained shape: {x_grid_pytorch.numpy().shape}, Expected shape: {x_grid.shape}") + print(f"test_meshgrid (Y) - Avg. Error: {error_y}, \n \ + Obtained shape: {y_grid_pytorch.numpy().shape}, Expected shape: {y_grid.shape}") + +if __name__ == '__main__': + + test_meshgrid() \ No newline at end of file diff --git a/cv/crestereo/function_convertion_tests/test_offset.py b/cv/crestereo/function_convertion_tests/test_offset.py new file mode 100644 index 0000000..fb8bf57 --- /dev/null +++ b/cv/crestereo/function_convertion_tests/test_offset.py @@ -0,0 +1,31 @@ +import pickle +import numpy as np +import megengine as mge + +import torch +import torch.nn.functional as F + +def test_offset(): + # Getting back the megengine objects: + with open('test_data/offset_test.pkl', 'rb') as f: + x_grid, y_grid, reshape_shape, transpose_order, expand_size, repeat_size, repeat_axis, offsets = pickle.load(f) + + x_grid = torch.tensor(x_grid.numpy()) + y_grid = torch.tensor(y_grid.numpy()) + offsets_mge = offsets.numpy() + N = repeat_size + + # Test Pytorch + offsets = torch.stack((x_grid, y_grid)) + offsets = offsets.reshape(2, -1).permute(1, 0) + for d in sorted((0, 2, 3)): + offsets = offsets.unsqueeze(d) + offsets = offsets.repeat_interleave(N, dim=0) + + error = np.mean(offsets.numpy()-offsets_mge) + print(f"test_offset - Avg. Error: {error}, \n \ + Obtained shape: {offsets.numpy().shape}, Expected shape: {offsets_mge.shape}") + +if __name__ == '__main__': + + test_offset() \ No newline at end of file diff --git a/cv/crestereo/function_convertion_tests/test_split.py b/cv/crestereo/function_convertion_tests/test_split.py new file mode 100644 index 0000000..7022eb8 --- /dev/null +++ b/cv/crestereo/function_convertion_tests/test_split.py @@ -0,0 +1,47 @@ +import pickle +import numpy as np +import megengine as mge + +import torch +import torch.nn.functional as F + +def test_split(): + # Getting back the megengine objects: + with open('test_data/split_test.pkl', 'rb') as f: + left_feature, size, axis, lefts = pickle.load(f) + + left_feature = torch.tensor(left_feature.numpy()) + + # Test Pytorch + lefts_pytorch = torch.split(left_feature, left_feature.shape[axis]//size, dim=axis) + + for i, (left_pytorch, left) in enumerate(zip(lefts_pytorch, lefts)): + + error = np.mean(left_pytorch.numpy()-left.numpy()) + print(f"test_split {i} - Avg. Error: {error}, \n \ + Obtained shape: {left_pytorch.numpy().shape}, Expected shape: {left.numpy().shape}\n") + +def test_split_list(): + # Getting back the megengine objects: + with open('test_data/split_test_list.pkl', 'rb') as f: + fmap1, size, axis, net, inp = pickle.load(f) + + fmap1 = torch.tensor(fmap1.numpy()) + net = net.numpy() + inp = inp.numpy() + + # Test Pytorch + net_pytorch, inp_pytorch = torch.split(fmap1, [size[0],size[0]], dim=axis) + + error_net = np.mean(net_pytorch.numpy()-net) + error_inp = np.mean(inp_pytorch.numpy()-inp) + print(f"test_split_list (net) - Avg. Error: {error_net}, \n \ + Obtained shape: {net_pytorch.numpy().shape}, Expected shape: {net.shape}\n") + print(f"test_split_list (inp) - Avg. Error: {error_inp}, \n \ + Obtained shape: {inp_pytorch.numpy().shape}, Expected shape: {inp.shape}\n") + + +if __name__ == '__main__': + + test_split() + test_split_list() \ No newline at end of file diff --git a/cv/crestereo/nets/__init__.py b/cv/crestereo/nets/__init__.py new file mode 100644 index 0000000..0d0730f --- /dev/null +++ b/cv/crestereo/nets/__init__.py @@ -0,0 +1 @@ +from .crestereo import CREStereo as Model diff --git a/cv/crestereo/nets/attention/__init__.py b/cv/crestereo/nets/attention/__init__.py new file mode 100644 index 0000000..a7f763c --- /dev/null +++ b/cv/crestereo/nets/attention/__init__.py @@ -0,0 +1,2 @@ +from .transformer import LocalFeatureTransformer +from .position_encoding import PositionEncodingSine diff --git a/cv/crestereo/nets/attention/linear_attention.py b/cv/crestereo/nets/attention/linear_attention.py new file mode 100644 index 0000000..61b1b85 --- /dev/null +++ b/cv/crestereo/nets/attention/linear_attention.py @@ -0,0 +1,81 @@ +""" +Linear Transformer proposed in "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention" +Modified from: https://github.com/idiap/fast-transformers/blob/master/fast_transformers/attention/linear_attention.py +""" + +import torch +from torch.nn import Module, Dropout + + +def elu_feature_map(x): + return torch.nn.functional.elu(x) + 1 + + +class LinearAttention(Module): + def __init__(self, eps=1e-6): + super().__init__() + self.feature_map = elu_feature_map + self.eps = eps + + def forward(self, queries, keys, values, q_mask=None, kv_mask=None): + """ Multi-Head linear attention proposed in "Transformers are RNNs" + Args: + queries: [N, L, H, D] + keys: [N, S, H, D] + values: [N, S, H, D] + q_mask: [N, L] + kv_mask: [N, S] + Returns: + queried_values: (N, L, H, D) + """ + Q = self.feature_map(queries) + K = self.feature_map(keys) + + # set padded position to zero + if q_mask is not None: + Q = Q * q_mask[:, :, None, None] + if kv_mask is not None: + K = K * kv_mask[:, :, None, None] + values = values * kv_mask[:, :, None, None] + + v_length = values.size(1) + values = values / v_length # prevent fp16 overflow + KV = torch.einsum("nshd,nshv->nhdv", K, values) # (S,D)' @ S,V + Z = 1 / (torch.einsum("nlhd,nhd->nlh", Q, K.sum(dim=1)) + self.eps) + queried_values = torch.einsum("nlhd,nhdv,nlh->nlhv", Q, KV, Z) * v_length + + return queried_values.contiguous() + + +class FullAttention(Module): + def __init__(self, use_dropout=False, attention_dropout=0.1): + super().__init__() + self.use_dropout = use_dropout + self.dropout = Dropout(attention_dropout) + + def forward(self, queries, keys, values, q_mask=None, kv_mask=None): + """ Multi-head scaled dot-product attention, a.k.a full attention. + Args: + queries: [N, L, H, D] + keys: [N, S, H, D] + values: [N, S, H, D] + q_mask: [N, L] + kv_mask: [N, S] + Returns: + queried_values: (N, L, H, D) + """ + + # Compute the unnormalized attention and apply the masks + QK = torch.einsum("nlhd,nshd->nlsh", queries, keys) + if kv_mask is not None: + QK.masked_fill_(~(q_mask[:, :, None, None] * kv_mask[:, None, :, None]), float('-inf')) + + # Compute the attention and the weighted average + softmax_temp = 1. / queries.size(3)**.5 # sqrt(D) + A = torch.softmax(softmax_temp * QK, dim=2) + if self.use_dropout: + A = self.dropout(A) + + queried_values = torch.einsum("nlsh,nshd->nlhd", A, values) + + return queried_values.contiguous() \ No newline at end of file diff --git a/cv/crestereo/nets/attention/position_encoding.py b/cv/crestereo/nets/attention/position_encoding.py new file mode 100644 index 0000000..c78307d --- /dev/null +++ b/cv/crestereo/nets/attention/position_encoding.py @@ -0,0 +1,41 @@ +import math +import torch +from torch import nn + + +class PositionEncodingSine(nn.Module): + """ + This is a sinusoidal position encoding that generalized to 2-dimensional images + """ + + def __init__(self, d_model, max_shape=(256, 256), temp_bug_fix=False): + """ + Args: + max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels + temp_bug_fix (bool): As noted in this [issue](https://github.com/zju3dv/LoFTR/issues/41), + the original implementation of LoFTR includes a bug in the pos-enc impl, which has little impact + on the final performance. For now, we keep both impls for backward compatability. + We will remove the buggy impl after re-training all variants of our released models. + """ + super().__init__() + pe = torch.zeros((d_model, *max_shape)) + y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0) + x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0) + if temp_bug_fix: + div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / (d_model//2))) + else: # a buggy implementation (for backward compatability only) + div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / d_model//2)) + div_term = div_term[:, None, None] # [C//4, 1, 1] + pe[0::4, :, :] = torch.sin(x_position * div_term) + pe[1::4, :, :] = torch.cos(x_position * div_term) + pe[2::4, :, :] = torch.sin(y_position * div_term) + pe[3::4, :, :] = torch.cos(y_position * div_term) + + self.register_buffer('pe', pe.unsqueeze(0), persistent=False) # [1, C, H, W] + + def forward(self, x): + """ + Args: + x: [N, C, H, W] + """ + return x + self.pe[:, :, :x.size(2), :x.size(3)].to(x.device) \ No newline at end of file diff --git a/cv/crestereo/nets/attention/transformer.py b/cv/crestereo/nets/attention/transformer.py new file mode 100644 index 0000000..de55ffc --- /dev/null +++ b/cv/crestereo/nets/attention/transformer.py @@ -0,0 +1,100 @@ +import copy +import torch +import torch.nn as nn +from .linear_attention import LinearAttention, FullAttention + +#Ref: https://github.com/zju3dv/LoFTR/blob/master/src/loftr/loftr_module/transformer.py +class LoFTREncoderLayer(nn.Module): + def __init__(self, + d_model, + nhead, + attention='linear'): + super(LoFTREncoderLayer, self).__init__() + + self.dim = d_model // nhead + self.nhead = nhead + + # multi-head attention + self.q_proj = nn.Linear(d_model, d_model, bias=False) + self.k_proj = nn.Linear(d_model, d_model, bias=False) + self.v_proj = nn.Linear(d_model, d_model, bias=False) + self.attention = LinearAttention() if attention == 'linear' else FullAttention() + self.merge = nn.Linear(d_model, d_model, bias=False) + + # feed-forward network + self.mlp = nn.Sequential( + nn.Linear(d_model*2, d_model*2, bias=False), + nn.ReLU(), + nn.Linear(d_model*2, d_model, bias=False), + ) + + # norm and dropout + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + + def forward(self, x, source, x_mask=None, source_mask=None): + """ + Args: + x (torch.Tensor): [N, L, C] + source (torch.Tensor): [N, S, C] + x_mask (torch.Tensor): [N, L] (optional) + source_mask (torch.Tensor): [N, S] (optional) + """ + bs = x.size(0) + query, key, value = x, source, source + + # multi-head attention + query = self.q_proj(query).view(bs, -1, self.nhead, self.dim) # [N, L, (H, D)] + key = self.k_proj(key).view(bs, -1, self.nhead, self.dim) # [N, S, (H, D)] + value = self.v_proj(value).view(bs, -1, self.nhead, self.dim) + message = self.attention(query, key, value, q_mask=x_mask, kv_mask=source_mask) # [N, L, (H, D)] + message = self.merge(message.view(bs, -1, self.nhead*self.dim)) # [N, L, C] + message = self.norm1(message) + + # feed-forward network + message = self.mlp(torch.cat([x, message], dim=2)) + message = self.norm2(message) + + return x + message + + +class LocalFeatureTransformer(nn.Module): + """A Local Feature Transformer (LoFTR) module.""" + + def __init__(self, d_model, nhead, layer_names, attention): + super(LocalFeatureTransformer, self).__init__() + + self.d_model = d_model + self.nhead = nhead + self.layer_names = layer_names + encoder_layer = LoFTREncoderLayer(d_model, nhead, attention) + self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(len(self.layer_names))]) + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def forward(self, feat0, feat1, mask0=None, mask1=None): + """ + Args: + feat0 (torch.Tensor): [N, L, C] + feat1 (torch.Tensor): [N, S, C] + mask0 (torch.Tensor): [N, L] (optional) + mask1 (torch.Tensor): [N, S] (optional) + """ + assert self.d_model == feat0.size(2), "the feature number of src and transformer must be equal" + + for layer, name in zip(self.layers, self.layer_names): + + if name == 'self': + feat0 = layer(feat0, feat0, mask0, mask0) + feat1 = layer(feat1, feat1, mask1, mask1) + elif name == 'cross': + feat0 = layer(feat0, feat1, mask0, mask1) + feat1 = layer(feat1, feat0, mask1, mask0) + else: + raise KeyError + + return feat0, feat1 \ No newline at end of file diff --git a/cv/crestereo/nets/corr.py b/cv/crestereo/nets/corr.py new file mode 100644 index 0000000..bf0e286 --- /dev/null +++ b/cv/crestereo/nets/corr.py @@ -0,0 +1,148 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .utils import bilinear_sampler, coords_grid, manual_pad + +class AGCL: + """ + Implementation of Adaptive Group Correlation Layer (AGCL). + """ + + def __init__(self, fmap1, fmap2, att=None): + self.fmap1 = fmap1 + self.fmap2 = fmap2 + + self.att = att + + self.coords = coords_grid(fmap1.shape[0], fmap1.shape[2], fmap1.shape[3], fmap1.device) + + def __call__(self, flow, extra_offset, small_patch=False, iter_mode=False): + if iter_mode: + corr = self.corr_iter(self.fmap1, self.fmap2, flow, small_patch) + else: + corr = self.corr_att_offset( + self.fmap1, self.fmap2, flow, extra_offset, small_patch + ) + return corr + + def get_correlation(self, left_feature, right_feature, psize=(3, 3), dilate=(1, 1)): + + N, C, H, W = left_feature.shape + + di_y, di_x = dilate[0], dilate[1] + pady, padx = psize[0] // 2 * di_y, psize[1] // 2 * di_x + + right_pad = manual_pad(right_feature, pady, padx) + + corr_list = [] + for h in range(0, pady * 2 + 1, di_y): + for w in range(0, padx * 2 + 1, di_x): + right_crop = right_pad[:, :, h : h + H, w : w + W] + assert right_crop.shape == left_feature.shape + corr = torch.mean(left_feature * right_crop, dim=1, keepdims=True) + corr_list.append(corr) + + corr_final = torch.cat(corr_list, dim=1) + + return corr_final + + def corr_iter(self, left_feature, right_feature, flow, small_patch): + + coords = self.coords + flow + coords = coords.permute(0, 2, 3, 1) + right_feature = bilinear_sampler(right_feature, coords) + + if small_patch: + psize_list = [(3, 3), (3, 3), (3, 3), (3, 3)] + dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)] + else: + psize_list = [(1, 9), (1, 9), (1, 9), (1, 9)] + dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)] + + N, C, H, W = left_feature.shape + lefts = torch.split(left_feature, left_feature.shape[1]//4, dim=1) + rights = torch.split(right_feature, right_feature.shape[1]//4, dim=1) + + corrs = [] + for i in range(len(psize_list)): + corr = self.get_correlation( + lefts[i], rights[i], psize_list[i], dilate_list[i] + ) + corrs.append(corr) + + final_corr = torch.cat(corrs, dim=1) + + return final_corr + + def corr_att_offset( + self, left_feature, right_feature, flow, extra_offset, small_patch + ): + + N, C, H, W = left_feature.shape + + if self.att is not None: + left_feature = left_feature.permute(0, 2, 3, 1).reshape(N, H * W, C) # 'n c h w -> n (h w) c' + right_feature = right_feature.permute(0, 2, 3, 1).reshape(N, H * W, C) # 'n c h w -> n (h w) c' + # 'n (h w) c -> n c h w' + left_feature, right_feature = self.att(left_feature, right_feature) + # 'n (h w) c -> n c h w' + left_feature, right_feature = [ + x.reshape(N, H, W, C).permute(0, 3, 1, 2) + for x in [left_feature, right_feature] + ] + + lefts = torch.split(left_feature, left_feature.shape[1]//4, dim=1) + rights = torch.split(right_feature, right_feature.shape[1]//4, dim=1) + + C = C // 4 + + if small_patch: + psize_list = [(3, 3), (3, 3), (3, 3), (3, 3)] + dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)] + else: + psize_list = [(1, 9), (1, 9), (1, 9), (1, 9)] + dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)] + + search_num = 9 + extra_offset = extra_offset.reshape(N, search_num, 2, H, W).permute(0, 1, 3, 4, 2) # [N, search_num, 1, 1, 2] + + corrs = [] + for i in range(len(psize_list)): + left_feature, right_feature = lefts[i], rights[i] + psize, dilate = psize_list[i], dilate_list[i] + + psizey, psizex = psize[0], psize[1] + dilatey, dilatex = dilate[0], dilate[1] + + ry = psizey // 2 * dilatey + rx = psizex // 2 * dilatex + x_grid, y_grid = torch.meshgrid(torch.arange(-rx, rx + 1, dilatex, device=self.fmap1.device), + torch.arange(-ry, ry + 1, dilatey, device=self.fmap1.device), indexing='xy') + + offsets = torch.stack((x_grid, y_grid)) + offsets = offsets.reshape(2, -1).permute(1, 0) + for d in sorted((0, 2, 3)): + offsets = offsets.unsqueeze(d) + offsets = offsets.repeat_interleave(N, dim=0) + offsets = offsets + extra_offset + + coords = self.coords + flow # [N, 2, H, W] + coords = coords.permute(0, 2, 3, 1) # [N, H, W, 2] + coords = torch.unsqueeze(coords, 1) + offsets + coords = coords.reshape(N, -1, W, 2) # [N, search_num*H, W, 2] + + right_feature = bilinear_sampler( + right_feature, coords + ) # [N, C, search_num*H, W] + right_feature = right_feature.reshape(N, C, -1, H, W) # [N, C, search_num, H, W] + left_feature = left_feature.unsqueeze(2).repeat_interleave(right_feature.shape[2], dim=2) + + corr = torch.mean(left_feature * right_feature, dim=1) + + corrs.append(corr) + + final_corr = torch.cat(corrs, dim=1) + + return final_corr diff --git a/cv/crestereo/nets/crestereo.py b/cv/crestereo/nets/crestereo.py new file mode 100644 index 0000000..3f99917 --- /dev/null +++ b/cv/crestereo/nets/crestereo.py @@ -0,0 +1,258 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .update import BasicUpdateBlock +from .extractor import BasicEncoder +from .corr import AGCL + +from .attention import PositionEncodingSine, LocalFeatureTransformer + +try: + autocast = torch.cuda.amp.autocast +except: + # dummy autocast for PyTorch < 1.6 + class autocast: + def __init__(self, enabled): + pass + def __enter__(self): + pass + def __exit__(self, *args): + pass + +#Ref: https://github.com/princeton-vl/RAFT/blob/master/core/raft.py +class CREStereo(nn.Module): + def __init__(self, max_disp=192, mixed_precision=False, test_mode=False): + super(CREStereo, self).__init__() + + self.max_flow = max_disp + self.mixed_precision = mixed_precision + self.test_mode = test_mode + + self.hidden_dim = 128 + self.context_dim = 128 + self.dropout = 0 + + self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=self.dropout) + self.update_block = BasicUpdateBlock(hidden_dim=self.hidden_dim, cor_planes=4 * 9, mask_size=4) + + # loftr + self.self_att_fn = LocalFeatureTransformer( + d_model=256, nhead=8, layer_names=["self"] * 1, attention="linear" + ) + self.cross_att_fn = LocalFeatureTransformer( + d_model=256, nhead=8, layer_names=["cross"] * 1, attention="linear" + ) + + # adaptive search + self.search_num = 9 + self.conv_offset_16 = nn.Conv2d( + 256, self.search_num * 2, kernel_size=3, stride=1, padding=1 + ) + self.conv_offset_8 = nn.Conv2d( + 256, self.search_num * 2, kernel_size=3, stride=1, padding=1 + ) + self.range_16 = 1 + self.range_8 = 1 + + def freeze_bn(self): + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + + def convex_upsample(self, flow, mask, rate=4): + """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """ + N, _, H, W = flow.shape + # print(flow.shape, mask.shape, rate) + mask = mask.view(N, 1, 9, rate, rate, H, W) + mask = torch.softmax(mask, dim=2) + + up_flow = F.unfold(rate * flow, [3,3], padding=1) + up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) + + up_flow = torch.sum(mask * up_flow, dim=2) + up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) + return up_flow.reshape(N, 2, rate*H, rate*W) + + def zero_init(self, fmap): + N, C, H, W = fmap.shape + _x = torch.zeros([N, 1, H, W], dtype=torch.float32) + _y = torch.zeros([N, 1, H, W], dtype=torch.float32) + zero_flow = torch.cat((_x, _y), dim=1).to(fmap.device) + return zero_flow + + def forward(self, image1, image2, flow_init=None, iters=10, upsample=True, test_mode=False): + """ Estimate optical flow between pair of frames """ + + image1 = 2 * (image1 / 255.0) - 1.0 + image2 = 2 * (image2 / 255.0) - 1.0 + + image1 = image1.contiguous() + image2 = image2.contiguous() + + hdim = self.hidden_dim + cdim = self.context_dim + + # run the feature network + with autocast(enabled=self.mixed_precision): + fmap1, fmap2 = self.fnet([image1, image2]) + + fmap1 = fmap1.float() + fmap2 = fmap2.float() + + with autocast(enabled=self.mixed_precision): + + # 1/4 -> 1/8 + # feature + fmap1_dw8 = F.avg_pool2d(fmap1, 2, stride=2) + fmap2_dw8 = F.avg_pool2d(fmap2, 2, stride=2) + + # offset + offset_dw8 = self.conv_offset_8(fmap1_dw8) + offset_dw8 = self.range_8 * (torch.sigmoid(offset_dw8) - 0.5) * 2.0 + + # context + net, inp = torch.split(fmap1, [hdim,hdim], dim=1) + net = torch.tanh(net) + inp = F.relu(inp) + net_dw8 = F.avg_pool2d(net, 2, stride=2) + inp_dw8 = F.avg_pool2d(inp, 2, stride=2) + + # 1/4 -> 1/16 + # feature + fmap1_dw16 = F.avg_pool2d(fmap1, 4, stride=4) + fmap2_dw16 = F.avg_pool2d(fmap2, 4, stride=4) + offset_dw16 = self.conv_offset_16(fmap1_dw16) + offset_dw16 = self.range_16 * (torch.sigmoid(offset_dw16) - 0.5) * 2.0 + + # context + net_dw16 = F.avg_pool2d(net, 4, stride=4) + inp_dw16 = F.avg_pool2d(inp, 4, stride=4) + + # positional encoding and self-attention + pos_encoding_fn_small = PositionEncodingSine( + d_model=256, max_shape=(image1.shape[2] // 16, image1.shape[3] // 16) + ) + # 'n c h w -> n (h w) c' + x_tmp = pos_encoding_fn_small(fmap1_dw16) + fmap1_dw16 = x_tmp.permute(0, 2, 3, 1).reshape(x_tmp.shape[0], x_tmp.shape[2] * x_tmp.shape[3], x_tmp.shape[1]) + # 'n c h w -> n (h w) c' + x_tmp = pos_encoding_fn_small(fmap2_dw16) + fmap2_dw16 = x_tmp.permute(0, 2, 3, 1).reshape(x_tmp.shape[0], x_tmp.shape[2] * x_tmp.shape[3], x_tmp.shape[1]) + + fmap1_dw16, fmap2_dw16 = self.self_att_fn(fmap1_dw16, fmap2_dw16) + fmap1_dw16, fmap2_dw16 = [ + x.reshape(x.shape[0], image1.shape[2] // 16, -1, x.shape[2]).permute(0, 3, 1, 2) + for x in [fmap1_dw16, fmap2_dw16] + ] + + corr_fn = AGCL(fmap1, fmap2) + corr_fn_dw8 = AGCL(fmap1_dw8, fmap2_dw8) + corr_fn_att_dw16 = AGCL(fmap1_dw16, fmap2_dw16, att=self.cross_att_fn) + + # Cascaded refinement (1/16 + 1/8 + 1/4) + predictions = [] + flow = None + flow_up = None + if flow_init is not None: + scale = fmap1.shape[2] / flow_init.shape[2] + flow = -scale * F.interpolate( + flow_init, + size=(fmap1.shape[2], fmap1.shape[3]), + mode="bilinear", + align_corners=True, + ) + else: + # zero initialization + flow_dw16 = self.zero_init(fmap1_dw16) + + # Recurrent Update Module + # RUM: 1/16 + for itr in range(iters // 2): + if itr % 2 == 0: + small_patch = False + else: + small_patch = True + + flow_dw16 = flow_dw16.detach() + out_corrs = corr_fn_att_dw16( + flow_dw16, offset_dw16, small_patch=small_patch + ) + + with autocast(enabled=self.mixed_precision): + net_dw16, up_mask, delta_flow = self.update_block( + net_dw16, inp_dw16, out_corrs, flow_dw16 + ) + + flow_dw16 = flow_dw16 + delta_flow + flow = self.convex_upsample(flow_dw16, up_mask, rate=4) + flow_up = -4 * F.interpolate( + flow, + size=(4 * flow.shape[2], 4 * flow.shape[3]), + mode="bilinear", + align_corners=True, + ) + predictions.append(flow_up) + + scale = fmap1_dw8.shape[2] / flow.shape[2] + flow_dw8 = -scale * F.interpolate( + flow, + size=(fmap1_dw8.shape[2], fmap1_dw8.shape[3]), + mode="bilinear", + align_corners=True, + ) + + # RUM: 1/8 + for itr in range(iters // 2): + if itr % 2 == 0: + small_patch = False + else: + small_patch = True + + flow_dw8 = flow_dw8.detach() + out_corrs = corr_fn_dw8(flow_dw8, offset_dw8, small_patch=small_patch) + + with autocast(enabled=self.mixed_precision): + net_dw8, up_mask, delta_flow = self.update_block( + net_dw8, inp_dw8, out_corrs, flow_dw8 + ) + + flow_dw8 = flow_dw8 + delta_flow + flow = self.convex_upsample(flow_dw8, up_mask, rate=4) + flow_up = -2 * F.interpolate( + flow, + size=(2 * flow.shape[2], 2 * flow.shape[3]), + mode="bilinear", + align_corners=True, + ) + predictions.append(flow_up) + + scale = fmap1.shape[2] / flow.shape[2] + flow = -scale * F.interpolate( + flow, + size=(fmap1.shape[2], fmap1.shape[3]), + mode="bilinear", + align_corners=True, + ) + + # RUM: 1/4 + for itr in range(iters): + if itr % 2 == 0: + small_patch = False + else: + small_patch = True + + flow = flow.detach() + out_corrs = corr_fn(flow, None, small_patch=small_patch, iter_mode=True) + + with autocast(enabled=self.mixed_precision): + net, up_mask, delta_flow = self.update_block(net, inp, out_corrs, flow) + + flow = flow + delta_flow + flow_up = -self.convex_upsample(flow, up_mask, rate=4) + predictions.append(flow_up) + + if self.test_mode: + return flow_up + + return predictions diff --git a/cv/crestereo/nets/extractor.py b/cv/crestereo/nets/extractor.py new file mode 100644 index 0000000..993cd3a --- /dev/null +++ b/cv/crestereo/nets/extractor.py @@ -0,0 +1,123 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +# Ref: https://github.com/princeton-vl/RAFT/blob/master/core/extractor.py +class ResidualBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(ResidualBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes) + self.norm2 = nn.BatchNorm2d(planes) + self.norm3 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes, affine=False) + self.norm2 = nn.InstanceNorm2d(planes, affine=False) + self.norm3 = nn.InstanceNorm2d(planes, affine=False) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + self.norm3 = nn.Sequential() + + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3) + + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + + x = self.downsample(x) + + return self.relu(x+y) + + +class BasicEncoder(nn.Module): + def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): + super(BasicEncoder, self).__init__() + self.norm_fn = norm_fn + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64, affine=False) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=2) + self.layer3 = self._make_layer(128, stride=1) + + # output convolution + self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) + + self.dropout = None + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x): + + # if input is list, combine batch dimension + is_list = isinstance(x, tuple) or isinstance(x, list) + if is_list: + batch_dim = x[0].shape[0] + x = torch.cat(x, dim=0) + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + + x = self.conv2(x) + + if self.dropout is not None: + x = self.dropout(x) + + if is_list: + x = torch.split(x, x.shape[0]//2, dim=0) + + return x \ No newline at end of file diff --git a/cv/crestereo/nets/update.py b/cv/crestereo/nets/update.py new file mode 100644 index 0000000..401d504 --- /dev/null +++ b/cv/crestereo/nets/update.py @@ -0,0 +1,91 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +#Ref: https://github.com/princeton-vl/RAFT/blob/master/core/update.py +class FlowHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256): + super(FlowHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class SepConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=192+128): + super(SepConvGRU, self).__init__() + self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + + self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + + def forward(self, h, x): + # horizontal + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + + # vertical + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + + return h + + +class BasicMotionEncoder(nn.Module): + def __init__(self, cor_planes): + super(BasicMotionEncoder, self).__init__() + + self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0) + self.convc2 = nn.Conv2d(256, 192, 3, padding=1) + self.convf1 = nn.Conv2d(2, 128, 7, padding=3) + self.convf2 = nn.Conv2d(128, 64, 3, padding=1) + self.conv = nn.Conv2d(64+192, 128-2, 3, padding=1) + + def forward(self, flow, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + flo = F.relu(self.convf1(flow)) + flo = F.relu(self.convf2(flo)) + + cor_flo = torch.cat([cor, flo], dim=1) + out = F.relu(self.conv(cor_flo)) + return torch.cat([out, flow], dim=1) + + +class BasicUpdateBlock(nn.Module): + def __init__(self, hidden_dim, cor_planes, mask_size=8): + super(BasicUpdateBlock, self).__init__() + + self.encoder = BasicMotionEncoder(cor_planes) + self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim) + self.flow_head = FlowHead(hidden_dim, hidden_dim=256) + + self.mask = nn.Sequential( + nn.Conv2d(128, 256, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(256, mask_size**2 *9, 1, padding=0)) + + def forward(self, net, inp, corr, flow, upsample=True): + # print(inp.shape, corr.shape, flow.shape) + motion_features = self.encoder(flow, corr) + # print(motion_features.shape, inp.shape) + inp = torch.cat((inp, motion_features), dim=1) + + net = self.gru(net, inp) + delta_flow = self.flow_head(net) + + # scale mask to balence gradients + mask = .25 * self.mask(net) + return net, mask, delta_flow diff --git a/cv/crestereo/nets/utils/__init__.py b/cv/crestereo/nets/utils/__init__.py new file mode 100644 index 0000000..8cc08e5 --- /dev/null +++ b/cv/crestereo/nets/utils/__init__.py @@ -0,0 +1 @@ +from .utils import bilinear_sampler, coords_grid, manual_pad \ No newline at end of file diff --git a/cv/crestereo/nets/utils/utils.py b/cv/crestereo/nets/utils/utils.py new file mode 100644 index 0000000..3a28205 --- /dev/null +++ b/cv/crestereo/nets/utils/utils.py @@ -0,0 +1,108 @@ +import torch +import torch.nn.functional as F +import numpy as np + +#Ref: https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + xgrid, ygrid = coords.split([1,1], dim=-1) + xgrid = 2*xgrid/(W-1) - 1 + ygrid = 2*ygrid/(H-1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + # img = F.grid_sample(img, grid, align_corners=True) + img = bilinear_grid_sample(img, grid, align_corners=True) + + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + +def coords_grid(batch, ht, wd, device): + coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device), indexing='ij') + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + +def manual_pad(x, pady, padx): + + pad = (padx, padx, pady, pady) + return F.pad(x.clone().detach(), pad, "replicate") + +# Ref: https://zenn.dev/pinto0309/scraps/7d4032067d0160 +def bilinear_grid_sample(im, grid, align_corners=False): + """Given an input and a flow-field grid, computes the output using input + values and pixel locations from grid. Supported only bilinear interpolation + method to sample the input pixels. + + Args: + im (torch.Tensor): Input feature map, shape (N, C, H, W) + grid (torch.Tensor): Point coordinates, shape (N, Hg, Wg, 2) + align_corners {bool}: If set to True, the extrema (-1 and 1) are + considered as referring to the center points of the input’s + corner pixels. If set to False, they are instead considered as + referring to the corner points of the input’s corner pixels, + making the sampling more resolution agnostic. + + Returns: + torch.Tensor: A tensor with sampled points, shape (N, C, Hg, Wg) + """ + n, c, h, w = im.shape + gn, gh, gw, _ = grid.shape + assert n == gn + + x = grid[:, :, :, 0] + y = grid[:, :, :, 1] + + if align_corners: + x = ((x + 1) / 2) * (w - 1) + y = ((y + 1) / 2) * (h - 1) + else: + x = ((x + 1) * w - 1) / 2 + y = ((y + 1) * h - 1) / 2 + + x = x.view(n, -1) + y = y.view(n, -1) + + x0 = torch.floor(x).long() + y0 = torch.floor(y).long() + x1 = x0 + 1 + y1 = y0 + 1 + + wa = ((x1 - x) * (y1 - y)).unsqueeze(1) + wb = ((x1 - x) * (y - y0)).unsqueeze(1) + wc = ((x - x0) * (y1 - y)).unsqueeze(1) + wd = ((x - x0) * (y - y0)).unsqueeze(1) + + # Apply default for grid_sample function zero padding + im_padded = torch.nn.functional.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0) + padded_h = h + 2 + padded_w = w + 2 + # save points positions after padding + x0, x1, y0, y1 = x0 + 1, x1 + 1, y0 + 1, y1 + 1 + + # Clip coordinates to padded image size + x0 = torch.where(x0 < 0, torch.tensor(0, device=im.device), x0) + x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x0) + x1 = torch.where(x1 < 0, torch.tensor(0, device=im.device), x1) + x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x1) + y0 = torch.where(y0 < 0, torch.tensor(0, device=im.device), y0) + y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y0) + y1 = torch.where(y1 < 0, torch.tensor(0, device=im.device), y1) + y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y1) + + im_padded = im_padded.view(n, c, -1) + + x0_y0 = (x0 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1) + x0_y1 = (x0 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1) + x1_y0 = (x1 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1) + x1_y1 = (x1 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1) + + Ia = torch.gather(im_padded, 2, x0_y0) + Ib = torch.gather(im_padded, 2, x0_y1) + Ic = torch.gather(im_padded, 2, x1_y0) + Id = torch.gather(im_padded, 2, x1_y1) + + return (Ia * wa + Ib * wb + Ic * wc + Id * wd).reshape(n, c, gh, gw) diff --git a/cv/crestereo/test_model.py b/cv/crestereo/test_model.py new file mode 100644 index 0000000..47b8c2a --- /dev/null +++ b/cv/crestereo/test_model.py @@ -0,0 +1,82 @@ +import torch +import torch.nn.functional as F +import numpy as np +import cv2 +from imread_from_url import imread_from_url + +from nets import Model + +device = 'cuda' + +#Ref: https://github.com/megvii-research/CREStereo/blob/master/test.py +def inference(left, right, model, n_iter=20): + + print("Model Forwarding...") + imgL = left.transpose(2, 0, 1) + imgR = right.transpose(2, 0, 1) + imgL = np.ascontiguousarray(imgL[None, :, :, :]) + imgR = np.ascontiguousarray(imgR[None, :, :, :]) + + imgL = torch.tensor(imgL.astype("float32")).to(device) + imgR = torch.tensor(imgR.astype("float32")).to(device) + + imgL_dw2 = F.interpolate( + imgL, + size=(imgL.shape[2] // 2, imgL.shape[3] // 2), + mode="bilinear", + align_corners=True, + ) + imgR_dw2 = F.interpolate( + imgR, + size=(imgL.shape[2] // 2, imgL.shape[3] // 2), + mode="bilinear", + align_corners=True, + ) + # print(imgR_dw2.shape) + with torch.inference_mode(): + pred_flow_dw2 = model(imgL_dw2, imgR_dw2, iters=n_iter, flow_init=None) + + pred_flow = model(imgL, imgR, iters=n_iter, flow_init=pred_flow_dw2) + pred_disp = torch.squeeze(pred_flow[:, 0, :, :]).cpu().detach().numpy() + + return pred_disp + +if __name__ == '__main__': + + left_img = imread_from_url("https://raw.githubusercontent.com/megvii-research/CREStereo/master/img/test/left.png") + right_img = imread_from_url("https://raw.githubusercontent.com/megvii-research/CREStereo/master/img/test/right.png") + + in_h, in_w = left_img.shape[:2] + + # Resize image in case the GPU memory overflows + eval_h, eval_w = (in_h,in_w) + assert eval_h%8 == 0, "input height should be divisible by 8" + assert eval_w%8 == 0, "input width should be divisible by 8" + + imgL = cv2.resize(left_img, (eval_w, eval_h), interpolation=cv2.INTER_LINEAR) + imgR = cv2.resize(right_img, (eval_w, eval_h), interpolation=cv2.INTER_LINEAR) + + model_path = "models/crestereo_eth3d.pth" + + model = Model(max_disp=256, mixed_precision=False, test_mode=True) + model.load_state_dict(torch.load(model_path), strict=True) + model.to(device) + model.eval() + + pred = inference(imgL, imgR, model, n_iter=20) + + t = float(in_w) / float(eval_w) + disp = cv2.resize(pred, (in_w, in_h), interpolation=cv2.INTER_LINEAR) * t + + disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0 + disp_vis = disp_vis.astype("uint8") + disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO) + + combined_img = np.hstack((left_img, disp_vis)) + cv2.namedWindow("output", cv2.WINDOW_NORMAL) + cv2.imshow("output", combined_img) + cv2.imwrite("output.jpg", disp_vis) + cv2.waitKey(0) + + + diff --git a/cv/crestereo/test_onnx_model.py b/cv/crestereo/test_onnx_model.py new file mode 100644 index 0000000..febf169 --- /dev/null +++ b/cv/crestereo/test_onnx_model.py @@ -0,0 +1,78 @@ +import numpy as np +import cv2 + +import onnxruntime + +# Ref: https://github.com/megvii-research/CREStereo/blob/master/test.py +def inference(left, right, model, no_flow_model): + # Get onnx model layer names (see convert_to_onnx.py for what these are) + input1_name = model.get_inputs()[0].name + input2_name = model.get_inputs()[1].name + input3_name = model.get_inputs()[2].name + output_name = model.get_outputs()[0].name + + # Decimate the image to half the original size for flow estimation network + imgL_dw2 = cv2.resize( + left, (left.shape[1] // 2, left.shape[0] // 2), interpolation=cv2.INTER_LINEAR) + imgR_dw2 = cv2.resize( + right, (right.shape[1] // 2, right.shape[0] // 2), interpolation=cv2.INTER_LINEAR) + + # Reshape inputs to match what is expected + imgL = left.transpose(2, 0, 1) + imgR = right.transpose(2, 0, 1) + imgL = np.ascontiguousarray(imgL[None, :, :, :]).astype("float32") + imgR = np.ascontiguousarray(imgR[None, :, :, :]).astype("float32") + + imgL_dw2 = imgL_dw2.transpose(2, 0, 1) + imgR_dw2 = imgR_dw2.transpose(2, 0, 1) + imgL_dw2 = np.ascontiguousarray(imgL_dw2[None, :, :, :]).astype("float32") + imgR_dw2 = np.ascontiguousarray(imgR_dw2[None, :, :, :]).astype("float32") + + print("Model Forwarding...") + # First pass it just to get the flow + pred_flow_dw2 = no_flow_model.run( + [output_name], {input1_name: imgL_dw2, input2_name: imgR_dw2})[0] + # Second pass gets us the disparity + pred_disp = model.run([output_name], { + input1_name: imgL, input2_name: imgR, input3_name: pred_flow_dw2})[0] + + return np.squeeze(pred_disp[:, 0, :, :]) + + +if __name__ == '__main__': + + left_img = cv2.imread("left.png") + right_img = cv2.imread("right.png") + + in_h, in_w = left_img.shape[:2] + + # Resize images + eval_h, eval_w = (in_h, in_w) + assert eval_h % 8 == 0, "input height should be divisible by 8" + assert eval_w % 8 == 0, "input width should be divisible by 8" + + imgL = cv2.resize(left_img, (eval_w, eval_h), + interpolation=cv2.INTER_LINEAR) + imgR = cv2.resize(right_img, (eval_w, eval_h), + interpolation=cv2.INTER_LINEAR) + + no_flow_model_path = "models/crestereo_without_flow.onnx" + model_path = "models/crestereo.onnx" + + model = onnxruntime.InferenceSession(model_path) + no_flow_model = onnxruntime.InferenceSession(no_flow_model_path) + + pred = inference(imgL, imgR, model, no_flow_model) + + t = float(in_w) / float(eval_w) + disp = cv2.resize(pred, (eval_w, eval_h), + interpolation=cv2.INTER_LINEAR) * t + disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0 + disp_vis = disp_vis.astype("uint8") + disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO) + + combined_img = np.hstack((left_img, disp_vis)) + cv2.namedWindow("output", cv2.WINDOW_NORMAL) + cv2.imshow("output", combined_img) + cv2.imwrite("output.jpg", disp_vis) + cv2.waitKey(0) diff --git a/cv/crestereo/train.py b/cv/crestereo/train.py new file mode 100644 index 0000000..1b7ecdc --- /dev/null +++ b/cv/crestereo/train.py @@ -0,0 +1,492 @@ +import argparse +import os +import shutil +import sys +import time +import logging +from collections import namedtuple +from itertools import repeat + +import yaml +from tensorboardX import SummaryWriter + +from nets import Model +from dataset import CREStereoDataset + +import torch +import torch.nn as nn +import torch.optim as optim +import torch.backends.cudnn as cudnn +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader, RandomSampler + + +def parse_yaml(file_path: str) -> namedtuple: + """Parse yaml configuration file and return the object in `namedtuple`.""" + with open(file_path, "rb") as f: + cfg: dict = yaml.safe_load(f) + args = namedtuple("train_args", cfg.keys())(*cfg.values()) + # save cfg into train_log + ensure_dir(args.log_dir) + dst_file = os.path.join(args.log_dir, file_path.split('/')[-1]) + shutil.copy2(file_path, dst_file) + return args + + +def format_time(elapse): + elapse = int(elapse) + hour = elapse // 3600 + minute = elapse % 3600 // 60 + seconds = elapse % 60 + return "{:02d}:{:02d}:{:02d}".format(hour, minute, seconds) + + +def ensure_dir(path): + if not os.path.exists(path): + os.makedirs(path, exist_ok=True) + + +def adjust_learning_rate(optimizer, epoch): + + warm_up = 0.02 + const_range = 0.6 + min_lr_rate = 0.05 + + if epoch <= args.n_total_epoch * warm_up: + lr = (1 - min_lr_rate) * args.base_lr / ( + args.n_total_epoch * warm_up + ) * epoch + min_lr_rate * args.base_lr + elif args.n_total_epoch * warm_up < epoch <= args.n_total_epoch * const_range: + lr = args.base_lr + else: + lr = (min_lr_rate - 1) * args.base_lr / ( + (1 - const_range) * args.n_total_epoch + ) * epoch + (1 - min_lr_rate * const_range) / (1 - const_range) * args.base_lr + + for param_group in optimizer.param_groups: + param_group['lr'] = lr + +def sequence_loss(flow_preds, flow_gt, valid, gamma=0.8): + ''' + valid: (2, 384, 512) (B, H, W) -> (B, 1, H, W) + flow_preds[0]: (B, 2, H, W) + flow_gt: (B, 2, H, W) + ''' + n_predictions = len(flow_preds) + flow_loss = 0.0 + for i in range(n_predictions): + i_weight = gamma ** (n_predictions - i - 1) + i_loss = torch.abs(flow_preds[i] - flow_gt) + flow_loss += i_weight * (valid.unsqueeze(1) * i_loss).mean() + + return flow_loss + +def repeater(data_loader): + for loader in repeat(data_loader): + for data in loader: + yield data + +def train_dist(args, world_size): + parser = argparse.ArgumentParser() + parser.add_argument("--local_rank",type=int) + FLAGS = parser.parse_args() + local_rank = FLAGS.local_rank + # directory check + log_model_dir = os.path.join(args.log_dir, "models") + ensure_dir(log_model_dir) + + # distributed init and model / optimizer + torch.cuda.set_device(local_rank) + dist.init_process_group(backend='nccl') # nccl is highly recommanded + model = Model( + max_disp=args.max_disp, mixed_precision=args.mixed_precision, test_mode=False + ) + # sync batch norm + model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(local_rank) + model = DDP(model, device_ids=[local_rank], output_device=local_rank) + optimizer = optim.Adam(model.parameters(), lr=0.1, betas=(0.9, 0.999)) + + if dist.get_rank() == 0: + # tensorboard + tb_log = SummaryWriter(os.path.join(args.log_dir, "train.events")) + + # worklog + logging.basicConfig(level=eval(args.log_level)) + worklog = logging.getLogger("train_logger") + worklog.propagate = False + fileHandler = logging.FileHandler( + os.path.join(args.log_dir, "worklog.txt"), mode="a", encoding="utf8" + ) + formatter = logging.Formatter( + fmt="%(asctime)s %(message)s", datefmt="%Y/%m/%d %H:%M:%S" + ) + fileHandler.setFormatter(formatter) + consoleHandler = logging.StreamHandler(sys.stdout) + formatter = logging.Formatter( + fmt="\x1b[32m%(asctime)s\x1b[0m %(message)s", datefmt="%Y/%m/%d %H:%M:%S" + ) + consoleHandler.setFormatter(formatter) + worklog.handlers = [fileHandler, consoleHandler] + + # params stat + worklog.info(f"Use {world_size} GPU(s)") + worklog.info("Params: %s" % sum([p.numel() for p in model.parameters()])) + + # load pretrained model if exist + chk_path = os.path.join(log_model_dir, "latest.pth") + if args.loadmodel is not None: + chk_path = args.loadmodel + elif not os.path.exists(chk_path): + chk_path = None + + if chk_path is not None: + if dist.get_rank() == 0: + worklog.info(f"loading model: {chk_path}") + # map_location=torch.device('cpu') make more balance memory usage + state_dict = torch.load(chk_path, map_location=torch.device('cpu')) + model.module.load_state_dict(state_dict['state_dict']) + optimizer.load_state_dict(state_dict['optim_state_dict']) + resume_epoch_idx = state_dict["epoch"] + resume_iters = state_dict["iters"] + start_epoch_idx = resume_epoch_idx + 1 + start_iters = resume_iters + else: + start_epoch_idx = 1 + start_iters = 0 + + # datasets + dataset = CREStereoDataset(args.training_data_path) + # dataset = MixDataset("train", + # data_path=args.data["train"]["data_path"], + # fields=args.data["train"]["fields"], + # filelists=args.data["train"]["filelists"], + # input_size=(args.data["train"]["input_size"][0], args.data["train"]["input_size"][1])) + if dist.get_rank() == 0: + worklog.info(f"Dataset size: {len(dataset)}") + train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) + dataloader = torch.utils.data.DataLoader(dataset, + batch_size=args.batch_size, num_workers=4, sampler=train_sampler) + + # counter + cur_iters = start_iters + total_iters = args.minibatch_per_epoch * args.n_total_epoch + t0 = time.perf_counter() + for epoch_idx in range(start_epoch_idx, args.n_total_epoch + 1): + dataloader.sampler.set_epoch(epoch_idx) + # adjust learning rate + epoch_total_train_loss = 0 + adjust_learning_rate(optimizer, epoch_idx) + model.train() + + t1 = time.perf_counter() + + for batch_idx, mini_batch_data in enumerate(dataloader): + + if batch_idx % args.minibatch_per_epoch == 0 and batch_idx != 0: + break + cur_iters += 1 + + # parse data + left, right, gt_disp, valid_mask = ( + mini_batch_data["left"].to(local_rank), + mini_batch_data["right"].to(local_rank), + mini_batch_data["disparity"].to(local_rank), + mini_batch_data["mask"].to(local_rank), + ) + + t2 = time.perf_counter() + optimizer.zero_grad() + + # pre-process + gt_disp = torch.unsqueeze(gt_disp, dim=1) # [2, 384, 512] -> [2, 1, 384, 512] + gt_flow = torch.cat([gt_disp, gt_disp * 0], dim=1) # [2, 2, 384, 512] + + # forward + flow_predictions = model(left, right) + + # loss & backword + loss = sequence_loss( + flow_predictions, gt_flow, valid_mask, gamma=0.8 + ).to(local_rank) + + # loss stats + loss_item = loss.data.item() + epoch_total_train_loss += loss_item + loss.backward() + optimizer.step() + t3 = time.perf_counter() + if dist.get_rank() == 0: + if cur_iters % 10 == 0: + tdata = t2 - t1 + time_train_passed = t3 - t0 + time_iter_passed = t3 - t1 + step_passed = cur_iters - start_iters + eta = ( + (total_iters - cur_iters) + / max(step_passed, 1e-7) + * time_train_passed + ) + + meta_info = list() + meta_info.append("{:.2g} b/s".format(1.0 / time_iter_passed)) + meta_info.append("passed:{}".format(format_time(time_train_passed))) + meta_info.append("eta:{}".format(format_time(eta))) + meta_info.append( + "data_time:{:.2g}".format(tdata / time_iter_passed) + ) + + meta_info.append( + "lr:{:.5g}".format(optimizer.param_groups[0]["lr"]) + ) + meta_info.append( + "[{}/{}:{}/{}]".format( + epoch_idx, + args.n_total_epoch, + batch_idx, + args.minibatch_per_epoch, + ) + ) + loss_info = list() + loss_info.append("{}:{:.4g}".format("total_loss", loss_item)) + # exp_name = ['\n' + os.path.basename(os.getcwd())] + + info = [",".join(meta_info+loss_info)] + worklog.info("".join(info)) + + # minibatch loss + tb_log.add_scalar("train/loss_batch", loss_item, cur_iters) + tb_log.add_scalar( + "train/lr", optimizer.param_groups[0]["lr"], cur_iters + ) + tb_log.flush() + + t1 = time.perf_counter() + + if dist.get_rank() == 0: + # epoch loss + tb_log.add_scalar( + "train/loss", + epoch_total_train_loss / args.minibatch_per_epoch, + epoch_idx, + ) + tb_log.flush() + + # save model params + ckp_data = { + "epoch": epoch_idx, + "iters": cur_iters, + "batch_size": args.batch_size * world_size, + "epoch_size": args.minibatch_per_epoch, + "train_loss": epoch_total_train_loss / args.minibatch_per_epoch, + "state_dict": model.module.state_dict(), + "optim_state_dict": optimizer.state_dict(), + } + torch.save(ckp_data, os.path.join(log_model_dir, "latest.pth")) + if epoch_idx % args.model_save_freq_epoch == 0: + save_path = os.path.join(log_model_dir, "epoch-%d.pth" % epoch_idx) + worklog.info(f"Model params saved: {save_path}") + torch.save(ckp_data, save_path) + if dist.get_rank() == 0: + worklog.info("Training is done, exit.") + +def train(args, world_size): + # directory check + log_model_dir = os.path.join(args.log_dir, "models") + ensure_dir(log_model_dir) + + # model / optimizer + model = Model( + max_disp=args.max_disp, mixed_precision=args.mixed_precision, test_mode=False + ) + model = nn.DataParallel(model,device_ids=[i for i in range(world_size)]) + model.cuda() + optimizer = optim.Adam(model.parameters(), lr=0.1, betas=(0.9, 0.999)) + + tb_log = SummaryWriter(os.path.join(args.log_dir, "train.events")) + + # worklog + logging.basicConfig(level=eval(args.log_level)) + worklog = logging.getLogger("train_logger") + worklog.propagate = False + fileHandler = logging.FileHandler( + os.path.join(args.log_dir, "worklog.txt"), mode="a", encoding="utf8" + ) + formatter = logging.Formatter( + fmt="%(asctime)s %(message)s", datefmt="%Y/%m/%d %H:%M:%S" + ) + fileHandler.setFormatter(formatter) + consoleHandler = logging.StreamHandler(sys.stdout) + formatter = logging.Formatter( + fmt="\x1b[32m%(asctime)s\x1b[0m %(message)s", datefmt="%Y/%m/%d %H:%M:%S" + ) + consoleHandler.setFormatter(formatter) + worklog.handlers = [fileHandler, consoleHandler] + + # params stat + worklog.info(f"Use {world_size} GPU(s)") + worklog.info("Params: %s" % sum([p.numel() for p in model.parameters()])) + + # load pretrained model if exist + chk_path = os.path.join(log_model_dir, "latest.pth") + if args.loadmodel is not None: + chk_path = args.loadmodel + elif not os.path.exists(chk_path): + chk_path = None + + if chk_path is not None: + worklog.info(f"loading model: {chk_path}") + state_dict = torch.load(chk_path) + model.module.load_state_dict(state_dict['state_dict']) + optimizer.load_state_dict(state_dict['optim_state_dict']) + resume_epoch_idx = state_dict["epoch"] + resume_iters = state_dict["iters"] + start_epoch_idx = resume_epoch_idx + 1 + start_iters = resume_iters + else: + start_epoch_idx = 1 + start_iters = 0 + + # datasets + dataset = CREStereoDataset(args.training_data_path) + sampler = RandomSampler(dataset, replacement=False) + worklog.info(f"Dataset size: {len(dataset)}") + dataloader = DataLoader(dataset, sampler=sampler, batch_size=args.batch_size*world_size, + num_workers=0, drop_last=True, persistent_workers=False, pin_memory=True) + dataloader = repeater(dataloader) + + # counter + cur_iters = start_iters + total_iters = args.minibatch_per_epoch * args.n_total_epoch + t0 = time.perf_counter() + for epoch_idx in range(start_epoch_idx, args.n_total_epoch + 1): + + # adjust learning rate + epoch_total_train_loss = 0 + adjust_learning_rate(optimizer, epoch_idx) + model.train() + + t1 = time.perf_counter() + + # for mini_batch_data in dataloader: + for batch_idx, mini_batch_data in enumerate(dataloader): + + if batch_idx % args.minibatch_per_epoch == 0 and batch_idx != 0: + break + cur_iters += 1 + + # parse data + left, right, gt_disp, valid_mask = ( + mini_batch_data["left"].cuda(), + mini_batch_data["right"].cuda(), + mini_batch_data["disparity"].cuda(), + mini_batch_data["mask"].cuda(), + ) + + t2 = time.perf_counter() + optimizer.zero_grad() + + # pre-process + gt_disp = torch.unsqueeze(gt_disp, dim=1) # [2, 384, 512] -> [2, 1, 384, 512] + gt_flow = torch.cat([gt_disp, gt_disp * 0], dim=1) # [2, 2, 384, 512] + + # forward + flow_predictions = model(left, right) + + # loss & backword + loss = sequence_loss( + flow_predictions, gt_flow, valid_mask, gamma=0.8 + ) + + # loss stats + loss_item = loss.data.item() + epoch_total_train_loss += loss_item + loss.backward() + optimizer.step() + t3 = time.perf_counter() + + if cur_iters % 10 == 0: + tdata = t2 - t1 + time_train_passed = t3 - t0 + time_iter_passed = t3 - t1 + step_passed = cur_iters - start_iters + eta = ( + (total_iters - cur_iters) + / max(step_passed, 1e-7) + * time_train_passed + ) + + meta_info = list() + meta_info.append("{:.2g} b/s".format(1.0 / time_iter_passed)) + meta_info.append("passed:{}".format(format_time(time_train_passed))) + meta_info.append("eta:{}".format(format_time(eta))) + meta_info.append( + "data_time:{:.2g}".format(tdata / time_iter_passed) + ) + + meta_info.append( + "lr:{:.5g}".format(optimizer.param_groups[0]["lr"]) + ) + meta_info.append( + "[{}/{}:{}/{}]".format( + epoch_idx, + args.n_total_epoch, + batch_idx, + args.minibatch_per_epoch, + ) + ) + loss_info = [" ==> {}:{:.4g}".format("loss", loss_item)] + # exp_name = ['\n' + os.path.basename(os.getcwd())] + + info = [",".join(meta_info)] + loss_info + worklog.info("".join(info)) + + # minibatch loss + tb_log.add_scalar("train/loss_batch", loss_item, cur_iters) + tb_log.add_scalar( + "train/lr", optimizer.param_groups[0]["lr"], cur_iters + ) + tb_log.flush() + + t1 = time.perf_counter() + + tb_log.add_scalar( + "train/loss", + epoch_total_train_loss / args.minibatch_per_epoch, + epoch_idx, + ) + tb_log.flush() + + # save model params + ckp_data = { + "epoch": epoch_idx, + "iters": cur_iters, + "batch_size": args.batch_size*world_size, + "epoch_size": args.minibatch_per_epoch, + "train_loss": epoch_total_train_loss / args.minibatch_per_epoch, + "state_dict": model.module.state_dict(), + "optim_state_dict": optimizer.state_dict(), + } + torch.save(ckp_data, os.path.join(log_model_dir, "latest.pth")) + if epoch_idx % args.model_save_freq_epoch == 0: + save_path = os.path.join(log_model_dir, "epoch-%d.pth" % epoch_idx) + worklog.info(f"Model params saved: {save_path}") + torch.save(ckp_data, save_path) + + worklog.info("Training is done, exit.") + +def main(args): + # initial info + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + world_size = torch.cuda.device_count() # number of GPU(s) + cudnn.benchmark = True + if args.dist and world_size > 1: + train_dist(args, world_size) + else: + train(args, world_size) + +if __name__ == "__main__": + # train configuration + args = parse_yaml("cfgs/train.yaml") + main(args) diff --git a/cv/defom-stereo/LICENSE b/cv/defom-stereo/LICENSE new file mode 100644 index 0000000..25169c9 --- /dev/null +++ b/cv/defom-stereo/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Insta360 Research Team + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/cv/defom-stereo/README.md b/cv/defom-stereo/README.md new file mode 100644 index 0000000..a2f6df7 --- /dev/null +++ b/cv/defom-stereo/README.md @@ -0,0 +1,225 @@ +# DEFOM-Stereo [CVPR 2025] + +The Official Pytorch Implementation for + +> [**DEFOM-Stereo: Depth Foundation Model Based Stereo Matching**](https://arxiv.org/abs/2501.09466) +> +> Authors: Hualie Jiang, Zhiqiang Lou, Laiyan Ding, Rui Xu, Minglang Tan, Wenjie Jiang and Rui Huang + +# Abstract +Stereo matching is a key technique for metric depth estimation in computer vision and robotics. +Real-world challenges like occlusion and non-texture hinder accurate disparity estimation from binocular matching cues. Recently, monocular relative depth estimation has shown remarkable generalization using vision foundation models. Thus, to facilitate robust stereo matching with monocular depth cues, we incorporate a robust monocular relative depth model into the recurrent stereo-matching framework, building a new framework for depth foundation model-based stereo-matching, DEFOM-Stereo. +In the feature extraction stage, we construct the combined context and matching feature encoder by integrating features from conventional CNNs and DEFOM. In the update stage, we use the depth predicted by DEFOM to initialize the recurrent disparity and introduce a scale update module to refine the disparity at the correct scale. DEFOM-Stereo is verified to have much stronger zero-shot generalization compared with SOTA methods. Moreover, DEFOM-Stereo achieves top performance on the KITTI 2012, KITTI 2015, Middlebury, and ETH3D benchmarks, ranking $1^{st}$ on many metrics. In the joint evaluation under the robust vision challenge, our model simultaneously outperforms previous models on the individual benchmarks, further demonstrating its outstanding capabilities. + + +# Pipeline +

+ +

+ +- We propose **a novel recurrent stereo-matching framework incorporating monocular depth cues** from a depth foundation model to improve robustness. +- We develop a simple technique that utilizes **pre-trained DEFOM features** to construct stronger **combined feature and context encoders**. +- We invent a **recurrent scale update** module empowered with the **scale lookup**, serving to recover accurate pixel-wise scales for the coarse DEFOM depth. + + + +# Zero-Shot Perfomance + +

+ +

+ +# Benchmark Performance + +

+ +

+ +# Robust Vision Challange + +

+ +

+ +# Preparation + +### Installation + +Create the environment + +```bash +conda env create -f environment.yaml +conda activate defomstereo +pip install -r requirements.txt +``` + + +### Datasets +The project requires the follow datasets: + + + + + + + + + + + + + + + + + + + + + + + +
KITTI-2012KITTI-2015MiddleburyETH3DInStereo2K
Virtual KITTI 2SceneFlowTartanAir +CREStereo Dataset +FallingThings +
Sintel StereoHR-VS3D Ken BurnsIRS DatasetBooster Dataset
+The datasets are organized as follows, + +``` +. +└── datasets + ├── 3dkenburns + │ ├── asdf-flying + │ ├── asdf-flying-depth + │ └── ... + ├── Booster_Dataset + │ ├── test + │ └── train + ├── CreStereo + │ ├── hole + │ ├── reflective + │ ├── shapenet + │ └── tree + ├── ETH3D + │ ├── two_view_testing + │ ├── two_view_training + │ └── two_view_training_gt + ├── FallingThings + │ └── fat + ├── HRVS + │ └── carla-highres + ├── InStereo2K + │ ├── part1 + │ ├── part2 + │ ├── part3 + │ ├── part4 + │ ├── part5 + │ └── test + ├── IRSDataset + │ ├── Home + │ ├── Office + │ ├── Restaurant + │ └── Store + ├── KITTI12 + │ ├── testing + │ └── training + ├── KITTI15 + │ ├── testing + │ └── training + ├── Middlebury + │ ├── 2005 + │ ├── 2006 + │ ├── 2014 + │ ├── 2021 + │ └── MiddEval3 + ├── SceneFlow + │ ├── Driving + │ ├── FlyingThings3D + │ └── Monkaa + ├── SintelStereo + │ └── training + ├── TartanAir + │ ├── abandonedfactory + │ ├── abandonedfactory_night + │ └── ... + └── VKITTI2 + ├── Scene01 + ├── Scene02 + ├── Scene06 + ├── Scene18 + └── Scene20 +``` + + +# Evaluation + +### Download the pre-trained models +``` +bash scripts/download_models.sh +``` + +The pretrained models are available on [Google Drive](https://drive.google.com/drive/folders/1cZLcIjHlmUo986gkR6FbofG1cj5BT36x?usp=sharing) and can be downloaded mamanually. + +### Perfom Evaluation +``` +bash scripts/evaluate.sh +``` + +# Make Benchmark Submission + +``` +bash scripts/make_submission.sh +``` + + +# Training + +### Download DaV2 models +``` +bash scripts/download_dav2.sh +``` + + +### Train on SceneFlow + +``` +bash scripts/train_sceneflow_vits.sh +bash scripts/train_sceneflow_vitl.sh +``` + +### Finetune for Benchmarks + +``` +bash scripts/train_kitti.sh +bash scripts/train_middlebury.sh +bash scripts/train_eth3d.sh +bash scripts/train_rvc.sh +``` + + +# Domo on real samples + +``` +python demo.py --restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth +``` + + + + +## Acknowledgements + +The project is based on [RAFT-Stereo](https://github.com/princeton-vl/RAFT-Stereo) and [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2) and we sincerely acknowledge their authors for opensourcing the excellent work. Besides, we would like to thank the CVPR reviewers and AC for their valuable feedback and recognition of our work. + + +## Citation + +Please cite our paper if you find our work useful in your research. + +``` +@inproceedings{jiang2025defom, + title={DEFOM-Stereo: Depth Foundation Model Based Stereo Matching}, + author={Jiang, Hualie and Lou, Zhiqiang and Ding, Laiyan and Xu, Rui and Tan, Minglang and Jiang, Wenjie and Huang, Rui}, + booktitle={IEEE International Conference on Computer Vision and Pattern Recognition (CVPR)}, + year={2025} +} +``` diff --git a/cv/defom-stereo/assets/benchmark.jpg b/cv/defom-stereo/assets/benchmark.jpg new file mode 100644 index 0000000..4fc4247 Binary files /dev/null and b/cv/defom-stereo/assets/benchmark.jpg differ diff --git a/cv/defom-stereo/assets/framework.svg b/cv/defom-stereo/assets/framework.svg new file mode 100644 index 0000000..0cabc54 --- /dev/null +++ b/cv/defom-stereo/assets/framework.svg @@ -0,0 +1,1357 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/cv/defom-stereo/assets/rvc.svg b/cv/defom-stereo/assets/rvc.svg new file mode 100644 index 0000000..8c8b280 --- /dev/null +++ b/cv/defom-stereo/assets/rvc.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/cv/defom-stereo/assets/zeroshot.jpg b/cv/defom-stereo/assets/zeroshot.jpg new file mode 100644 index 0000000..4b8b5e3 Binary files /dev/null and b/cv/defom-stereo/assets/zeroshot.jpg differ diff --git a/cv/defom-stereo/core/__init__.py b/cv/defom-stereo/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cv/defom-stereo/core/corr.py b/cv/defom-stereo/core/corr.py new file mode 100644 index 0000000..37d7ec9 --- /dev/null +++ b/cv/defom-stereo/core/corr.py @@ -0,0 +1,212 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from core.utils.utils import bilinear_sampler + +try: + import corr_sampler +except: + pass + +try: + import alt_cuda_corr +except: + # alt_cuda_corr is not compiled + pass + +class CorrSampler(torch.autograd.Function): + @staticmethod + def forward(ctx, volume, coords, radius): + ctx.save_for_backward(volume,coords) + ctx.radius = radius + corr, = corr_sampler.forward(volume, coords, radius) + return corr + @staticmethod + def backward(ctx, grad_output): + volume, coords = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_volume, = corr_sampler.backward(volume, coords, grad_output, ctx.radius) + return grad_volume, None, None + + +class CorrBlockFast1D: + def __init__(self, fmap1, fmap2, num_levels=4, radius=4, **kwargs): + self.num_levels = num_levels + self.radius = radius + self.corr_pyramid = [] + # all pairs correlation + corr = CorrBlockFast1D.corr(fmap1, fmap2) + batch, h1, w1, dim, w2 = corr.shape + corr = corr.reshape(batch*h1*w1, dim, 1, w2) + for i in range(self.num_levels): + self.corr_pyramid.append(corr.view(batch, h1, w1, -1, w2//2**i)) + corr = F.avg_pool2d(corr, [1, 2], stride=[1, 2]) + + def __call__(self, coords): + out_pyramid = [] + bz, _, ht, wd = coords.shape + coords = coords[:, [0]] + for i in range(self.num_levels): + corr = CorrSampler.apply(self.corr_pyramid[i].squeeze(3), coords/2**i, self.radius) + out_pyramid.append(corr.view(bz, -1, ht, wd)) + return torch.cat(out_pyramid, dim=1) + + @staticmethod + def corr(fmap1, fmap2): + B, D, H, W1 = fmap1.shape + _, _, _, W2 = fmap2.shape + fmap1 = fmap1.view(B, D, H, W1) + fmap2 = fmap2.view(B, D, H, W2) + corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2) + corr = corr.reshape(B, H, W1, 1, W2).contiguous() + return corr / torch.sqrt(torch.tensor(D).float()) + + +class PytorchAlternateCorrBlock1D: + def __init__(self, fmap1, fmap2, num_levels=4, radius=4, **kwargs): + self.num_levels = num_levels + self.radius = radius + self.corr_pyramid = [] + self.fmap1 = fmap1 + self.fmap2 = fmap2 + + def corr(self, fmap1, fmap2, coords): + B, D, H, W = fmap2.shape + # map grid coordinates to [-1,1] + xgrid, ygrid = coords.split([1,1], dim=-1) + xgrid = 2*xgrid/(W-1) - 1 + ygrid = 2*ygrid/(H-1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + output_corr = [] + for grid_slice in grid.unbind(3): + fmapw_mini = F.grid_sample(fmap2, grid_slice, align_corners=True) + corr = torch.sum(fmapw_mini * fmap1, dim=1) + output_corr.append(corr) + corr = torch.stack(output_corr, dim=1).permute(0,2,3,1) + + return corr / torch.sqrt(torch.tensor(D).float()) + + def __call__(self, coords): + r = self.radius + coords = coords.permute(0, 2, 3, 1) + batch, h1, w1, _ = coords.shape + fmap1 = self.fmap1 + fmap2 = self.fmap2 + out_pyramid = [] + for i in range(self.num_levels): + dx = torch.zeros(1) + dy = torch.linspace(-r, r, 2*r+1) + delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device) + centroid_lvl = coords.reshape(batch, h1, w1, 1, 2).clone() + centroid_lvl[..., 0] = centroid_lvl[..., 0] / 2**i + coords_lvl = centroid_lvl + delta.view(-1, 2) + corr = self.corr(fmap1, fmap2, coords_lvl) + fmap2 = F.avg_pool2d(fmap2, [1, 2], stride=[1, 2]) + out_pyramid.append(corr) + out = torch.cat(out_pyramid, dim=-1) + return out.permute(0, 3, 1, 2).contiguous().float() + + +class CorrBlock1D: + def __init__(self, fmap1, fmap2, coords, num_levels=4, radius=4, + scale_list=[0.25, 0.5, 2.0, 4.0], scale_corr_radius=4): + self.num_levels = num_levels + self.radius = radius + self.scale_list = scale_list + self.scale_corr_radius = scale_corr_radius + self.corr_pyramid = [] + self.coords_pyramid = [] + dx = torch.linspace(-radius, radius, 2*radius+1) + self.dx = dx[:, None].to(coords.device) + + sdx = torch.linspace(-scale_corr_radius, scale_corr_radius, 2*scale_corr_radius+1) + self.sdx = sdx[:, None].to(coords.device) + + # all pairs correlation + corr = CorrBlock1D.corr(fmap1, fmap2) + + batch, h1, w1, _, w2 = corr.shape + self.batch = batch + self.h1 = h1 + self.w1 = w1 + self.w2 = w2 + corr = corr.reshape(batch*h1*w1, 1, 1, w2) + self.coords = coords.reshape(batch*h1*w1, 1, 1, 1) + + self.corr_pyramid.append(corr) + for i in range(1, self.num_levels): + corr = F.avg_pool2d(corr, [1, 2], stride=[1, 2]) + self.corr_pyramid.append(corr) + + def __call__(self, disp, scaling=False): + batch, _, h1, w1 = disp.shape + + disp = disp.reshape(self.batch*self.h1*self.w1, 1, 1, 1) + out_pyramid = [] + + if scaling: + corr = self.corr_pyramid[0] + for scale in self.scale_list: + x0 = self.sdx + self.coords - scale * disp + y0 = torch.zeros_like(x0) + coords_lvl = torch.cat([x0, y0], dim=-1) + corr_s = bilinear_sampler(corr, coords_lvl) + corr_s = corr_s.view(self.batch, self.h1, self.w1, -1) + out_pyramid.append(corr_s) + else: + coords = self.coords - disp + for i in range(self.num_levels): + corr = self.corr_pyramid[i] + x0 = self.dx + coords / 2**i + y0 = torch.zeros_like(x0) + coords_lvl = torch.cat([x0, y0], dim=-1) + corr_s = bilinear_sampler(corr, coords_lvl) + corr_s = corr_s.view(self.batch, self.h1, self.w1, -1) + out_pyramid.append(corr_s) + + out = torch.cat(out_pyramid, dim=-1) + return out.permute(0, 3, 1, 2).contiguous().float() + + @staticmethod + def corr(fmap1, fmap2): + B, D, H, W1 = fmap1.shape + _, _, _, W2 = fmap2.shape + fmap1 = fmap1.view(B, D, H, W1) + fmap2 = fmap2.view(B, D, H, W2) + corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2) + corr = corr.reshape(B, H, W1, 1, W2).contiguous() + return corr / torch.sqrt(torch.tensor(D).float()) + + +class AlternateCorrBlock: + def __init__(self, fmap1, fmap2, num_levels=4, radius=4, **kwargs): + raise NotImplementedError + self.num_levels = num_levels + self.radius = radius + + self.pyramid = [(fmap1, fmap2)] + for i in range(1, self.num_levels): + fmap1 = F.avg_pool2d(fmap1, 2, stride=2) + fmap2 = F.avg_pool2d(fmap2, 2, stride=2) + self.pyramid.append((fmap1, fmap2)) + + def __call__(self, coords): + coords = coords.permute(0, 2, 3, 1) + B, H, W, _ = coords.shape + dim = self.pyramid[0][0].shape[1] + + corr_list = [] + for i in range(self.num_levels): + r = self.radius + fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous() + fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous() + + coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous() + corr, = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r) + corr_list.append(corr.squeeze(1)) + + corr = torch.stack(corr_list, dim=1) + corr = corr.reshape(B, -1, H, W) + return corr / torch.sqrt(torch.tensor(dim).float()) diff --git a/cv/defom-stereo/core/defom_stereo.py b/cv/defom-stereo/core/defom_stereo.py new file mode 100644 index 0000000..4d32a53 --- /dev/null +++ b/cv/defom-stereo/core/defom_stereo.py @@ -0,0 +1,142 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from core.update import BasicMultiUpdateBlock, ScaleBasicMultiUpdateBlock +from core.extractor import BasicEncoder, MultiBasicEncoder, ResidualBlock, DefomEncoder +from core.corr import CorrBlock1D, PytorchAlternateCorrBlock1D, CorrBlockFast1D, AlternateCorrBlock +from core.utils.utils import coords_grid, upflow, get_danv2_io_size + + +try: + autocast = torch.cuda.amp.autocast +except: + # dummy autocast for PyTorch < 1.6 + class autocast: + def __init__(self, enabled): + pass + def __enter__(self): + pass + def __exit__(self, *args): + pass + + +class DEFOMStereo(nn.Module): + def __init__(self, args): + super(DEFOMStereo, self).__init__() + self.args = args + + self.register_buffer('mean', torch.tensor([[0.485, 0.456, 0.406]])[..., None, None] * 255) + self.register_buffer('std', torch.tensor([[0.229, 0.224, 0.225]])[..., None, None] * 255) + + self.defomencoder = DefomEncoder(args.dinov2_encoder, idepth_scale=args.idepth_scale) + + context_dims = args.hidden_dims + + self.fnet = BasicEncoder(self.defomencoder.out_dim, output_dim=256, norm_fn='instance', downsample=args.n_downsample) + + self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, 3, padding=3//2) for i in range(self.args.n_gru_layers)]) + + self.update_block = BasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims) + self.scale_update_block = ScaleBasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims) + + self.cnet = MultiBasicEncoder(self.defomencoder.out_dim, output_dim=[args.hidden_dims, context_dims], + norm_fn=args.context_norm, downsample=args.n_downsample) + + def freeze_bn(self): + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + + def initialize_coords(self, img): + """ Disparity is represented as difference between two vertical coordinate grids disp + = coords0[:, :1] - coords1[:, :1] """ + N, _, H, W = img.shape + + coords = coords_grid(N, H, W)[:, :1].to(img.device) + + return coords + + def upsample_flow(self, flow, mask): + """ Upsample disparity field [H/scale, W/scale, 1] -> [H, W, 1] using convex combination """ + N, D, H, W = flow.shape + factor = 2 ** self.args.n_downsample + mask = mask.view(N, 1, 9, factor, factor, H, W) + mask = torch.softmax(mask, dim=2) + + up_flow = F.unfold(factor * flow, [3, 3], padding=1) + up_flow = up_flow.view(N, D, 9, 1, 1, H, W) + + up_flow = torch.sum(mask * up_flow, dim=2) + up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) + return up_flow.reshape(N, D, factor * H, factor * W) + + def forward(self, image1, image2, iters=12, scale_iters=3, test_mode=False): + """ Estimate optical flow between pair of frames """ + + image1 = ((image1 - self.mean)/self.std).contiguous() + image2 = ((image2 - self.mean)/self.std).contiguous() + + bs, _, h, w = image1.shape + danv2_io_sizes = get_danv2_io_size(h, w, self.args.n_downsample) + + # run the context network + with autocast(enabled=self.args.mixed_precision): + d_features, dfeat1, dfeat2, disp = self.defomencoder([image1, image2], danv2_io_sizes) + + cnet_list = self.cnet(image1, d_features) + fmap1, fmap2 = self.fnet([image1, image2], [dfeat1, dfeat2]) + net_list = [torch.tanh(x[0]) for x in cnet_list] + inp_list = [torch.relu(x[1]) for x in cnet_list] + # Rather than running the GRU's conv layers on the context features multiple times, we do it once at the beginning + inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i, conv in zip(inp_list, self.context_zqr_convs)] + + coords = self.initialize_coords(net_list[0]) + + fmap1, fmap2 = fmap1.float(), fmap2.float() + disp = disp.float() + corr_fn = CorrBlock1D(fmap1, fmap2, coords, radius=self.args.corr_radius, num_levels=self.args.corr_levels, + scale_list=self.args.scale_list, scale_corr_radius=self.args.scale_corr_radius) + + disp_predictions = [] + for itr in range(iters): + disp = disp.detach() + + if itr < scale_iters: + corr = corr_fn(disp, scaling=True) # index correlation volume + with autocast(enabled=self.args.mixed_precision): + net_list, up_mask, scale_disp = self.scale_update_block(net_list, inp_list, corr, disp, + iter32=self.args.n_gru_layers == 3, + iter16=self.args.n_gru_layers >= 2) + + # F(t+1) = \Scale(t) x F(t) + disp = scale_disp * disp + else: + corr = corr_fn(disp, scaling=False) # index correlation volume + with autocast(enabled=self.args.mixed_precision): + net_list, up_mask, delta_disp = self.update_block(net_list, inp_list, corr, disp, + iter32=self.args.n_gru_layers == 3, + iter16=self.args.n_gru_layers >= 2) + + # To avoid unstability, we limit the disparity update within the searching range. + delta_disp = torch.clip(delta_disp, min=-2**(self.args.corr_levels-1)*self.args.corr_radius, + max=2**(self.args.corr_levels-1)*self.args.corr_radius) + + # F(t+1) = F(t) + \Delta(t) + disp = disp + delta_disp + + # We do not need to upsample or output intermediate results in test_mode + if test_mode and itr < iters - 1: + continue + + # upsample predictions + if up_mask is None: + disp_up = upflow(disp, factor=2 ** self.n_downsample) + else: + disp_up = self.upsample_flow(disp, up_mask) + + disp_predictions.append(disp_up) + + if test_mode: + return disp_up + + return disp_predictions diff --git a/cv/defom-stereo/core/extractor.py b/cv/defom-stereo/core/extractor.py new file mode 100644 index 0000000..285b018 --- /dev/null +++ b/cv/defom-stereo/core/extractor.py @@ -0,0 +1,388 @@ +import os +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.models.layers import DropPath + +from depth_anything_v2.dpt import DepthAnythingV2 + + +class ConvBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(ConvBlock, self).__init__() + + self.conv = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes) + self.norm2 = nn.BatchNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes) + self.norm2 = nn.InstanceNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.Sequential() + + def forward(self, x): + + return self.relu(self.norm1(self.conv(x))) + + +class ResidualBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(ResidualBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes) + self.norm2 = nn.BatchNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes) + self.norm2 = nn.InstanceNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.Sequential() + + if stride == 1 and in_planes == planes: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3) + + def forward(self, x): + y = x + y = self.conv1(y) + y = self.norm1(y) + y = self.relu(y) + y = self.conv2(y) + y = self.norm2(y) + y = self.relu(y) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x+y) + + +class BottleneckBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1, ratio=4): + super(BottleneckBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes // ratio, kernel_size=1, padding=0) + self.conv2 = nn.Conv2d(planes // ratio, planes // ratio, kernel_size=3, padding=1, stride=stride) + self.conv3 = nn.Conv2d(planes // ratio, planes, kernel_size=1, padding=0) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // ratio) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // ratio) + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not (stride == 1 and in_planes == planes): + self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes // ratio) + self.norm2 = nn.BatchNorm2d(planes // ratio) + self.norm3 = nn.BatchNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm4 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes // ratio) + self.norm2 = nn.InstanceNorm2d(planes // ratio) + self.norm3 = nn.InstanceNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm4 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + self.norm3 = nn.Sequential() + if not (stride == 1 and in_planes == planes): + self.norm4 = nn.Sequential() + + if stride == 1 and in_planes == planes: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + y = self.relu(self.norm3(self.conv3(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BasicEncoder(nn.Module): + def __init__(self, d_dim, output_dim=128, norm_fn='batch', downsample=3): + super(BasicEncoder, self).__init__() + self.norm_fn = norm_fn + self.downsample = downsample + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=1 + (downsample > 1)) + self.layer3 = self._make_layer(128, stride=1 + (downsample > 0)) + + # depth feat convolution + self.convd = ConvBlock(d_dim, 128, self.norm_fn) + + # output convolution + self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x, dfeats): + + # if input is list, combine batch dimension + is_list = isinstance(x, tuple) or isinstance(x, list) + if is_list: + batch_dim = x[0].shape[0] + x = torch.cat(x, dim=0) + + is_list = isinstance(dfeats, tuple) or isinstance(dfeats, list) + if is_list: + batch_dim = dfeats[0].shape[0] + dfeats = torch.cat(dfeats, dim=0) + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + + x = x + self.convd(dfeats) + + x = self.conv2(x) + + if is_list: + x = x.split(split_size=batch_dim, dim=0) + + return x + + +class MultiBasicEncoder(nn.Module): + def __init__(self, d_dim, output_dim=[128, 128, 128], norm_fn='batch', downsample=3, drop_path_rate=0.2): + super(MultiBasicEncoder, self).__init__() + self.d_dim = d_dim + self.norm_fn = norm_fn + self.downsample = downsample + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=1 + (downsample > 1)) + self.layer3 = self._make_layer(128, stride=1 + (downsample > 0)) + self.layer4 = self._make_layer(128, stride=2) + self.layer5 = self._make_layer(128, stride=2) + + self.drop_path = DropPath(drop_path_rate) + + self.conv08 = ConvBlock(d_dim, 128, self.norm_fn) + output_list = [] + for dim in output_dim: + conv_out = nn.Sequential( + ResidualBlock(128, 128, self.norm_fn, stride=1), + nn.Conv2d(128, dim[2], 3, padding=1)) + output_list.append(conv_out) + + self.outputs08 = nn.ModuleList(output_list) + + self.conv16 = ConvBlock(d_dim, 128, self.norm_fn) + output_list = [] + for dim in output_dim: + conv_out = nn.Sequential( + ResidualBlock(128, 128, self.norm_fn, stride=1), + nn.Conv2d(128, dim[1], 3, padding=1)) + output_list.append(conv_out) + + self.outputs16 = nn.ModuleList(output_list) + + self.conv32 = ConvBlock(d_dim, 128, self.norm_fn) + output_list = [] + for dim in output_dim: + conv_out = nn.Conv2d(128, dim[0], 3, padding=1) + output_list.append(conv_out) + + self.outputs32 = nn.ModuleList(output_list) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x, d_feats, num_layers=3): + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + + feat = x + self.drop_path(self.conv08(d_feats[0])) + outputs08 = [f(feat) for f in self.outputs08] + if num_layers == 1: + return (outputs08,) + + y = self.layer4(x) + feat = y + self.drop_path(self.conv16(d_feats[1])) + outputs16 = [f(feat) for f in self.outputs16] + + if num_layers == 2: + return (outputs08, outputs16) + + z = self.layer5(y) + feat = z + self.drop_path(self.conv32(d_feats[2])) + outputs32 = [f(feat) for f in self.outputs32] + + return (outputs08, outputs16, outputs32) + + +class DefomEncoder(nn.Module): + def __init__(self, dinov2_encoder, pretrained=True, freeze=True, idepth_scale=0.25): + super(DefomEncoder, self).__init__() + self.dinov2_encoder = dinov2_encoder + self.idepth_scale = idepth_scale + self.pretrained = pretrained + self.freeze = freeze + + model_configs = { + 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}, + 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, + 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}, + 'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]} + } + + self.depth_anything = DepthAnythingV2(**model_configs[self.dinov2_encoder]) + + if pretrained and os.path.exists(f'./checkpoints/depth_anything_v2_{dinov2_encoder}.pth'): + self.depth_anything.load_state_dict( + torch.load(f'./checkpoints/depth_anything_v2_{dinov2_encoder}.pth', map_location='cpu'), strict=False) + if freeze: + for param in self.depth_anything.pretrained.parameters(): + param.requires_grad = False + for param in self.depth_anything.depth_head.parameters(): + param.requires_grad = False + + self.out_dim = model_configs[self.dinov2_encoder]['features'] + + def forward(self, x, danv2_io_sizes): + + x = torch.cat(x, dim=0) + ih, iw, oh, ow = danv2_io_sizes + x = F.interpolate(x, (ih, iw), mode="bilinear", align_corners=True) + + features, left_feat, right_feat, idepth = self.depth_anything(x, oh, ow) + + bs = idepth.shape[0] + max_idepth, _ = torch.max(idepth.view(bs, -1), dim=1) + max_idepth = max_idepth.detach().view(bs, 1, 1, 1) + 1e-8 + idepth = idepth / max_idepth * self.idepth_scale * ow + 0.01 + + return features, left_feat, right_feat, idepth diff --git a/cv/defom-stereo/core/stereo_datasets.py b/cv/defom-stereo/core/stereo_datasets.py new file mode 100644 index 0000000..64e96c9 --- /dev/null +++ b/cv/defom-stereo/core/stereo_datasets.py @@ -0,0 +1,583 @@ +# Data loading based on https://github.com/NVIDIA/flownet2-pytorch + +import numpy as np +from numpy import linalg as LA +import torch +import torch.utils.data as data +import torch.nn.functional as F +import logging +import os +import re +import copy +import math +import random +from pathlib import Path +from glob import glob +import os.path as osp + +from core.utils import frame_utils +from core.utils.augmentor import DispAugmentor, SparseDispAugmentor + + +class StereoDataset(data.Dataset): + def __init__(self, aug_params=None, sparse=False, reader=None, is_eval=False, is_test=False): + self.augmentor = None + self.sparse = sparse + if aug_params is not None and "crop_size" in aug_params: + if sparse: + self.augmentor = SparseDispAugmentor(**aug_params) + else: + self.augmentor = DispAugmentor(**aug_params) + + if reader is None: + self.disparity_reader = frame_utils.read_gen + else: + self.disparity_reader = reader + + self.is_eval = is_eval + self.is_test = is_test + self.init_seed = False + self.disparity_list = [] + self.image_list = [] + + # number of copies of the datasets + self.v = 1 + + def __getitem__(self, index): + + if self.is_test: + img1 = frame_utils.read_gen(self.image_list[index][0]) + img2 = frame_utils.read_gen(self.image_list[index][1]) + img1 = np.array(img1).astype(np.uint8) + img2 = np.array(img2).astype(np.uint8) + if len(img1.shape) == 2: + img1 = np.tile(img1[..., None], (1, 1, 3)) + img2 = np.tile(img2[..., None], (1, 1, 3)) + else: + img1 = img1[..., :3] + img2 = img2[..., :3] + img1 = torch.from_numpy(img1).permute(2, 0, 1).float() + img2 = torch.from_numpy(img2).permute(2, 0, 1).float() + return img1, img2, self.image_list[index][0] + + if not self.init_seed: + worker_info = torch.utils.data.get_worker_info() + if worker_info is not None: + torch.manual_seed(worker_info.id) + np.random.seed(worker_info.id) + random.seed(worker_info.id) + self.init_seed = True + + index = index % (len(self.image_list)*self.v) + index = index % len(self.image_list) + + if not self.is_eval and len(self.disparity_list[index]) > 1 and np.random.rand() > 0.5: + disp = self.disparity_reader(self.disparity_list[index][1]) + if isinstance(disp, tuple): + disp, valid = disp + else: + valid = disp < 1024 + img1 = frame_utils.read_gen(self.image_list[index][1]) + img2 = frame_utils.read_gen(self.image_list[index][0]) + + img1 = np.array(img1).astype(np.uint8)[:, ::-1] + img2 = np.array(img2).astype(np.uint8)[:, ::-1] + disp = np.array(disp).astype(np.float32)[:, ::-1] + valid = np.array(valid).astype(np.bool_)[:, ::-1] + + else: + disp = self.disparity_reader(self.disparity_list[index][0]) + if isinstance(disp, tuple): + disp, valid = disp + else: + valid = disp < 1024 + + img1 = frame_utils.read_gen(self.image_list[index][0]) + img2 = frame_utils.read_gen(self.image_list[index][1]) + + img1 = np.array(img1).astype(np.uint8) + img2 = np.array(img2).astype(np.uint8) + disp = np.array(disp).astype(np.float32) + valid = np.array(valid).astype(np.bool_) + + # grayscale images + if len(img1.shape) == 2: + img1 = np.tile(img1[..., None], (1, 1, 3)) + img2 = np.tile(img2[..., None], (1, 1, 3)) + else: + img1 = img1[..., :3] + img2 = img2[..., :3] + + if self.augmentor is not None: + if self.sparse: + img1, img2, disp, valid = self.augmentor(img1, img2, disp, valid) + else: + img1, img2, disp = self.augmentor(img1, img2, disp) + + img1 = torch.from_numpy(img1.copy()).permute(2, 0, 1).float() + img2 = torch.from_numpy(img2.copy()).permute(2, 0, 1).float() + disp = torch.from_numpy(disp[..., np.newaxis].copy()).permute(2, 0, 1).float() + if self.sparse: + valid = torch.from_numpy(valid[..., np.newaxis].astype(np.bool_).copy()).permute(2, 0, 1) + else: + valid = disp < 512 + + return {"img1": img1, "img2": img2, "disp": disp, "valid": valid, "imageL_file": self.image_list[index][0], "disp_file": self.disparity_list[index][0]} + + def __mul__(self, v): + self.v = v + return self + + def __len__(self): + return len(self.image_list)*self.v + + +class SceneFlowDatasets(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/SceneFlow/', dstype='frames_cleanpass', things_test=False): + super(SceneFlowDatasets, self).__init__(aug_params, is_eval=things_test) + self.root = root + self.dstype = dstype + + if things_test: + self._add_things("TEST") + else: + self._add_things("TRAIN") + self._add_monkaa() + self._add_driving() + + def _add_things(self, split='TRAIN'): + """ Add FlyingThings3D data """ + + original_length = len(self.disparity_list) + root = osp.join(self.root, 'FlyingThings3D') + left_images = sorted(glob(osp.join(root, self.dstype, split, '*/*/left/*.png'))) + right_images = [im.replace('left', 'right') for im in left_images] + disparity_images = [im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images] + + # Choose a random subset of 400 images for validation + state = np.random.get_state() + np.random.seed(1000) + val_idxs = set(np.random.permutation(len(left_images))[:400]) + np.random.set_state(state) + + for idx, (img1, img2, disp) in enumerate(zip(left_images, right_images, disparity_images)): + if (split == 'TEST' and idx in val_idxs) or split == 'TRAIN': + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + logging.info(f"Added {len(self.disparity_list) - original_length} from FlyingThings {self.dstype}") + + def _add_monkaa(self): + """ Add FlyingThings3D data """ + + original_length = len(self.disparity_list) + root = osp.join(self.root, 'Monkaa') + left_images = sorted(glob(osp.join(root, self.dstype, '*/left/*.png')) ) + right_images = [image_file.replace('left', 'right') for image_file in left_images ] + disparity_images = [im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ] + + for img1, img2, disp in zip(left_images, right_images, disparity_images): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + logging.info(f"Added {len(self.disparity_list) - original_length} from Monkaa {self.dstype}") + + def _add_driving(self): + """ Add FlyingThings3D data """ + + original_length = len(self.disparity_list) + root = osp.join(self.root, 'Driving') + left_images = sorted(glob(osp.join(root, self.dstype, '*/*/*/left/*.png')) ) + right_images = [image_file.replace('left', 'right') for image_file in left_images ] + disparity_images = [im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ] + + for img1, img2, disp in zip(left_images, right_images, disparity_images): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + logging.info(f"Added {len(self.disparity_list) - original_length} from Driving {self.dstype}") + + +class ETH3D(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/ETH3D', split='training', is_eval=False, is_test=False): + super(ETH3D, self).__init__(aug_params, sparse=True, is_eval=is_eval, is_test=is_test) + + image1_list = sorted(glob(osp.join(root, f'two_view_{split}/*/im0.png'))) + image2_list = sorted(glob(osp.join(root, f'two_view_{split}/*/im1.png'))) + disp_list = sorted(glob(osp.join(root, 'two_view_training_gt/*/disp0GT.pfm'))) if split == 'training'\ + else [osp.join(root, 'two_view_training_gt/playground_1l/disp0GT.pfm')]*len(image1_list) + + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp]] + + +class KITTI(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/KITTI', split='15', image_set='training', is_eval=False, is_test=False): + super(KITTI, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispKITTI, is_eval=is_eval, is_test=is_test) + assert split in ["12", "15"] + root = root + split + assert os.path.exists(root) + + if split == '15': + image1_list = sorted(glob(os.path.join(root, image_set, 'image_2/*_10.png'))) + image2_list = sorted(glob(os.path.join(root, image_set, 'image_3/*_10.png'))) + disp_list = sorted( + glob(os.path.join(root, 'training', 'disp_occ_0/*_10.png'))) if image_set == 'training' else [osp.join( + root, 'training/disp_occ_0/000085_10.png')]*len(image1_list) + else: + image1_list = sorted(glob(os.path.join(root, image_set, 'colored_0/*_10.png'))) + image2_list = sorted(glob(os.path.join(root, image_set, 'colored_1/*_10.png'))) + disp_list = sorted( + glob(os.path.join(root, 'training', 'disp_occ/*_10.png'))) if image_set == 'training' else [osp.join( + root, 'training/disp_occ/000085_10.png')] * len(image1_list) + + for idx, (img1, img2, disp) in enumerate(zip(image1_list, image2_list, disp_list)): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp]] + + +class Middlebury(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/Middlebury', split='F', image_set='training', is_eval=False, is_test=False): + super(Middlebury, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispMiddlebury, is_eval=is_eval, is_test=is_test) + assert os.path.exists(root) + assert split in ["F", "H", "Q", "2005", "2006", "2014", "2021"] + assert image_set in ["training", "test"] + + if split == "2005": + scenes = list((Path(root) / "2005").glob("*")) + for scene in scenes: + self.image_list += [[str(scene / "view1.png"), str(scene / "view5.png")]] + self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]] + for illum in ["1", "2", "3"]: + for exp in ["0", "1", "2"]: + self.image_list += [[str(scene / f"Illum{illum}/Exp{exp}/view1.png"), str(scene / f"Illum{illum}/Exp{exp}/view5.png")]] + self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]] + elif split == "2006": + scenes = list((Path(root) / "2006").glob("*")) + for scene in scenes: + self.image_list += [[str(scene / "view1.png"), str(scene / "view5.png")]] + self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]] + for illum in ["1", "2", "3"]: + for exp in ["0", "1", "2"]: + self.image_list += [[str(scene / f"Illum{illum}/Exp{exp}/view1.png"), str(scene / f"Illum{illum}/Exp{exp}/view5.png")]] + self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]] + elif split == "2014": # datasets/Middlebury/2014/Pipes-perfect/im0.png + scenes = list((Path(root) / "2014").glob("*")) + for scene in scenes: + for s in ["E", "L", ""]: + self.image_list += [[str(scene / "im0.png"), str(scene / f"im1{s}.png")]] + self.disparity_list += [[str(scene / "disp0.pfm"), str(scene / "disp1.pfm")]] + elif split == "2021": + scenes = list((Path(root) / "2021/data").glob("*")) + for scene in scenes: + self.image_list += [[str(scene / "im0.png"), str(scene / "im1.png")]] + self.disparity_list += [[str(scene / "disp0.pfm"), str(scene / "disp1.pfm")]] + for s in ["0", "1", "2", "3"]: + if os.path.exists(str(scene / f"ambient/L0/im0e{s}.png")): + self.image_list += [[str(scene / f"ambient/L0/im0e{s}.png"), str(scene / f"ambient/L0/im1e{s}.png")]] + self.disparity_list += [[str(scene / "disp0.pfm"), str(scene / "disp1.pfm")]] + else: + if image_set == 'training': + lines = list(map(osp.basename, glob(os.path.join(root, "MiddEval3/trainingF/*")))) + if is_eval: + lines = list(filter(lambda p: any(s in p.split('/') for s in Path(os.path.join(root, "MiddEval3/official_train.txt")).read_text().splitlines()), lines)) + else: + lines = list(map(osp.basename, glob(os.path.join(root, "MiddEval3/testF/*")))) + + image1_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/im0.png') for name in lines]) + image2_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/im1.png') for name in lines]) + + disp_list = sorted([os.path.join(root, "MiddEval3", f'training{split}', f'{name}/disp0GT.pfm') for name in lines]) \ + if image_set == 'training' else [os.path.join(root, "MiddEval3", f'training{split}', 'Adirondack/disp0GT.pfm')]*len(image1_list) + + assert len(image1_list) == len(image2_list) == len(disp_list) > 0, [image1_list, split] + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp]] + + +class SintelStereo(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/SintelStereo'): + super().__init__(aug_params, reader=frame_utils.readDispSintelStereo) + + image1_list = sorted(glob(osp.join(root, 'training/*_left/*/frame_*.png'))) + image2_list = sorted(glob(osp.join(root, 'training/*_right/*/frame_*.png'))) + disp_list = sorted(glob(osp.join(root, 'training/disparities/*/frame_*.png'))) * 2 + + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + assert img1.split('/')[-2:] == disp.split('/')[-2:] + self.image_list += [[img1, img2]] + self.disparity_list += [[disp]] + + +class FallingThings(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/FallingThings'): + super().__init__(aug_params, reader=frame_utils.readDispFallingThings) + assert os.path.exists(root) + + image1_list = sorted(glob(osp.join(root, 'fat/single/*/*/*.left.jpg'))) + \ + sorted(glob(osp.join(root, 'fat/mixed/*/*.left.jpg'))) + image2_list = [e.replace('left.jpg', 'right.jpg') for e in image1_list] + disp_list = [e.replace('left.jpg', 'left.depth.png') for e in image1_list] + + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + + +class TartanAir(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/TartanAir'): + super().__init__(aug_params, reader=frame_utils.readDispTartanAir) + assert os.path.exists(root) + + image1_list = sorted(glob(osp.join(root, '*/*/*/*/image_left/*_left.png'))) + image2_list = [e.replace('_left', '_right') for e in image1_list] + disp_list = [e.replace('image_left', 'depth_left').replace('left.png', 'left_depth.npy') for e in image1_list] + + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + + +class CarlaHighres(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/HRVS/carla-highres'): + super().__init__(aug_params) + assert os.path.exists(root) + + image1_list = sorted(glob(osp.join(root, 'trainingF/*/im0.png'))) + image2_list = [e.replace('im0', 'im1') for e in image1_list] + disp1_list = [e.replace('im0.png', 'disp0GT.pfm') for e in image1_list] + disp2_list = [e.replace('im1.png', 'disp1GT.pfm') for e in image2_list] + + for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp1, disp2]] + + +class InStereo2K(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/InStereo2K', split='training'): + super(InStereo2K, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispInStereo2K, is_eval=split!="training") + if split == "training": + image1_list = sorted(glob(osp.join(root, 'part*/*/left.png'))) + else: + image1_list = sorted(glob(osp.join(root, 'test/*/left.png'))) + + image2_list = [e.replace('left', 'right') for e in image1_list] + disp_list = [e.replace('left', 'left_disp') for e in image1_list] + + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + + +class CreStereo(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/CreStereo'): + super(CreStereo, self).__init__(aug_params, reader=frame_utils.readDispCreStereo) + + image1_list = sorted(glob(osp.join(root, '*/*_left.jpg'))) + image2_list = [e.replace('left', 'right') for e in image1_list] + disp_list = [e.replace('_left.jpg', '_left.disp.png') for e in image1_list] + + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp, disp.replace('left', 'right')]] + + +class IRS(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/IRSDataset'): + super().__init__(aug_params) + image1_list = sorted(glob(osp.join(root, '*/*/l_*.png'))) + image2_list = sorted(glob(osp.join(root, '*/*/r_*.png'))) + disp_list = sorted(glob(osp.join(root, '*/*/d_*.pfm'))) + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + assert img1.split('/')[-2] == disp.split('/')[-2] + assert img1.split('.')[0].split('_')[-1] == disp.split('.')[0].split('_')[-1] + if 'QAOfficeAndSecurityRoom2_Night' in img1: # bad scenes + continue + self.image_list += [[img1, img2]] + self.disparity_list += [[disp]] + + +class Booster(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/Booster_Dataset', split='train', is_eval=False, is_test=False): + super().__init__(aug_params, sparse=True, reader=frame_utils.readDispBooster, is_eval=is_eval, is_test=is_test) + assert os.path.exists(root) + + folder_list = sorted(glob(osp.join(root, split+'/balanced/*'))) + for folder in folder_list: + image1_list = sorted(glob(osp.join(folder, 'camera_00/im*.png'))) + image2_list = sorted(glob(osp.join(folder, 'camera_02/im*.png'))) + if split=="train": + for img1 in image1_list: + for img2 in image2_list: + self.image_list += [[img1, img2]] + self.disparity_list += [[osp.join(folder, 'disp_00.npy'), osp.join(folder, 'disp_02.npy')]] + else: + for img1, img2 in zip(image1_list, image2_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[osp.join(folder, 'disp_00.npy'), osp.join(folder, 'disp_02.npy')]] + + +class ThreeDKenBurns(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/3dkenburns'): + super().__init__(aug_params, reader=frame_utils.readDisp3DKenBurns) + + image1_list = sorted(glob(osp.join(root, '*/*l-image.png'))) + image2_list = sorted(glob(osp.join(root, '*/*r-image.png'))) + + disp1_list = sorted(glob(osp.join(root, '*/*l-depth.exr'))) + disp2_list = sorted(glob(osp.join(root, '*/*r-depth.exr'))) + + for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp1, disp2]] + + +class VKITTI2(StereoDataset): + def __init__(self, aug_params=None, root='./datasets/VKITTI2'): + super().__init__(aug_params, reader=frame_utils.readDispVKITTI2) + + image1_list = sorted(glob(osp.join(root, 'Scene*/*/frames/rgb/Camera_0/rgb_*.jpg'))) + image2_list = sorted(glob(osp.join(root, 'Scene*/*/frames/rgb/Camera_1/rgb_*.jpg'))) + + disp1_list = sorted(glob(osp.join(root, 'Scene*/*/frames/depth/Camera_0/depth_*.png'))) + disp2_list = sorted(glob(osp.join(root, 'Scene*/*/frames/depth/Camera_1/depth_*.png'))) + + for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list): + self.image_list += [[img1, img2]] + self.disparity_list += [[disp1, disp2]] + + +def fetch_dataloader(args): + """ Create the data loader for the corresponding trainign set """ + + aug_params = {'crop_size': args.image_size, 'min_scale': args.spatial_scale[0], + 'max_scale': args.spatial_scale[1], 'do_flip': False, 'yjitter': not args.noyjitter} + if hasattr(args, "saturation_range") and args.saturation_range is not None: + aug_params["saturation_range"] = args.saturation_range + if hasattr(args, "img_gamma") and args.img_gamma is not None: + aug_params["gamma"] = args.img_gamma + if hasattr(args, "do_flip") and args.do_flip is not None: + aug_params["do_flip"] = args.do_flip + + assert len(args.train_datasets) == len(args.train_folds) + + train_dataset = None + for fold, dataset_name in zip(args.train_folds, args.train_datasets): + if dataset_name.startswith("middlebury_"): + new_dataset = Middlebury(aug_params, split=dataset_name.replace('middlebury_','')) * fold + elif dataset_name == 'sceneflow': + clean_dataset = SceneFlowDatasets(aug_params, dstype='frames_cleanpass') + final_dataset = SceneFlowDatasets(aug_params, dstype='frames_finalpass') + new_dataset = clean_dataset*fold+final_dataset*fold + logging.info(f"Adding {len(new_dataset)} samples from SceneFlow") + elif 'kitti1' in dataset_name: + new_dataset = KITTI(aug_params, split=dataset_name[-2:], image_set='training') * fold + logging.info(f"Adding {len(new_dataset)} samples from KITTI"+dataset_name[-2:]) + elif 'eth3d' in dataset_name: + new_dataset = ETH3D(aug_params, split='training') * fold + logging.info(f"Adding {len(new_dataset)} samples from ETH3D") + elif dataset_name == 'sintel_stereo': + new_dataset = SintelStereo(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Sintel Stereo") + elif dataset_name == 'falling_things': + new_dataset = FallingThings(aug_params)*fold + logging.info(f"Adding {len(new_dataset)} samples from FallingThings") + elif dataset_name.startswith('tartan_air'): + new_dataset = TartanAir(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Tartain Air") + elif dataset_name.startswith('carla_highres'): + new_dataset = CarlaHighres(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Carla Highres") + elif dataset_name.startswith('irs'): + new_dataset = IRS(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from IRS") + elif dataset_name.startswith('crestereo'): + new_dataset = CreStereo(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from CreStereo") + elif dataset_name.startswith('instereo2k'): + new_dataset = InStereo2K(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from InStereo2K") + elif dataset_name.startswith('booster'): + new_dataset = Booster(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Booster") + elif dataset_name.startswith('3dkenburns'): + new_dataset = ThreeDKenBurns(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from 3D Ken Burns") + elif dataset_name.startswith('vkitti2'): + new_dataset = VKITTI2(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from VKITTI2") + + train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset + + train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, + pin_memory=True, shuffle=True, num_workers=int(os.environ.get('SLURM_CPUS_PER_TASK', 6))-2, drop_last=True) + + logging.info('Training with %d image pairs' % len(train_dataset)) + return train_loader + + +def fetch_dataset(args): + """ Create the dataset for the corresponding training set """ + + aug_params = {'crop_size': args.image_size, 'min_scale': args.spatial_scale[0], + 'max_scale': args.spatial_scale[1], 'do_flip': False, 'yjitter': not args.noyjitter} + if hasattr(args, "saturation_range") and args.saturation_range is not None: + aug_params["saturation_range"] = args.saturation_range + if hasattr(args, "img_gamma") and args.img_gamma is not None: + aug_params["gamma"] = args.img_gamma + if hasattr(args, "do_flip") and args.do_flip is not None: + aug_params["do_flip"] = args.do_flip + + assert len(args.train_datasets) == len(args.train_folds) + + train_dataset = None + for fold, dataset_name in zip(args.train_folds, args.train_datasets): + if dataset_name.startswith("middlebury_"): + new_dataset = Middlebury(aug_params, split=dataset_name.replace('middlebury_', '')) * fold + logging.info(f"Adding {len(new_dataset)} samples from {dataset_name}") + elif 'eth3d' in dataset_name: + new_dataset = ETH3D(aug_params, split='training') * fold + logging.info(f"Adding {len(new_dataset)} samples from ETH3D") + elif 'kitti1' in dataset_name: + new_dataset = KITTI(aug_params, split=dataset_name[-2:], image_set='training') * fold + logging.info(f"Adding {len(new_dataset)} samples from KITTI"+dataset_name[-2:]) + elif dataset_name == 'sceneflow': + clean_dataset = SceneFlowDatasets(aug_params, dstype='frames_cleanpass') + final_dataset = SceneFlowDatasets(aug_params, dstype='frames_finalpass') + new_dataset = clean_dataset*fold+final_dataset*fold + logging.info(f"Adding {len(new_dataset)} samples from SceneFlow") + elif dataset_name == 'sintel_stereo': + new_dataset = SintelStereo(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Sintel Stereo") + elif dataset_name == 'falling_things': + new_dataset = FallingThings(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from FallingThings") + elif dataset_name.startswith('tartan_air'): + new_dataset = TartanAir(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Tartain Air") + elif dataset_name.startswith('carla_highres'): + new_dataset = CarlaHighres(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Carla Highres") + elif dataset_name.startswith('irs'): + new_dataset = IRS(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from IRS") + elif dataset_name.startswith('crestereo'): + new_dataset = CreStereo(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from CreStereo") + elif dataset_name.startswith('instereo2k'): + new_dataset = InStereo2K(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from InStereo2K") + elif dataset_name.startswith('booster'): + new_dataset = Booster(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from Booster") + elif dataset_name.startswith('3dkenburns'): + new_dataset = ThreeDKenBurns(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from 3D Ken Burns") + elif dataset_name.startswith('vkitti2'): + new_dataset = VKITTI2(aug_params) * fold + logging.info(f"Adding {len(new_dataset)} samples from VKITTI2") + + train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset + + logging.info('Training with %d image pairs' % len(train_dataset)) + return train_dataset + diff --git a/cv/defom-stereo/core/update.py b/cv/defom-stereo/core/update.py new file mode 100644 index 0000000..47a9d99 --- /dev/null +++ b/cv/defom-stereo/core/update.py @@ -0,0 +1,195 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from opt_einsum import contract + + +class DispHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256, output_dim=1): + super(DispHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, output_dim, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class ConvGRU(nn.Module): + def __init__(self, hidden_dim, input_dim, kernel_size=3): + super(ConvGRU, self).__init__() + self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, + padding=kernel_size//2) + self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, + padding=kernel_size//2) + self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, + padding=kernel_size//2) + + def forward(self, h, cz, cr, cq, *x_list): + x = torch.cat(x_list, dim=1) + hx = torch.cat([h, x], dim=1) + + z = torch.sigmoid(self.convz(hx) + cz) + r = torch.sigmoid(self.convr(hx) + cr) + q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)) + cq) + + h = (1-z) * h + z * q + return h + + +class SepConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=192+128): + super(SepConvGRU, self).__init__() + self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + + self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + + def forward(self, h, *x): + # horizontal + x = torch.cat(x, dim=1) + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + + # vertical + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + + return h + + +class BasicMotionEncoder(nn.Module): + def __init__(self, cor_planes, c1_planes=64, c2_planes=64, f1_planes=64, f2_planes=64, out_planes=128): + super(BasicMotionEncoder, self).__init__() + + self.convc1 = nn.Conv2d(cor_planes, c1_planes, 1, padding=0) + self.convc2 = nn.Conv2d(c1_planes, c2_planes, 3, padding=1) + self.convd1 = nn.Conv2d(1, f1_planes, 7, padding=3) + self.convd2 = nn.Conv2d(f1_planes, f2_planes, 3, padding=1) + self.conv = nn.Conv2d(c2_planes+f2_planes, out_planes-1, 3, padding=1) + + def forward(self, disp, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + dis = F.relu(self.convd1(disp)) + dis = F.relu(self.convd2(dis)) + + cor_dis = torch.cat([cor, dis], dim=1) + out = F.relu(self.conv(cor_dis)) + return torch.cat([out, disp], dim=1) + + +def pool2x(x): + return F.avg_pool2d(x, 3, stride=2, padding=1) + + +def pool4x(x): + return F.avg_pool2d(x, 5, stride=4, padding=1) + + +def interp(x, dest): + interp_args = {'mode': 'bilinear', 'align_corners': True} + return F.interpolate(x, dest.shape[2:], **interp_args) + + +# for RAFT-Stereo +class BasicMultiUpdateBlock(nn.Module): + def __init__(self, args, hidden_dims=[128, 128, 128]): + super().__init__() + self.args = args + encoder_output_dim = 128 + cor_planes = args.corr_levels * (2*args.corr_radius + 1) + self.encoder = BasicMotionEncoder(cor_planes, out_planes=encoder_output_dim) + + self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.n_gru_layers > 1)) + self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.n_gru_layers == 3) + hidden_dims[2]) + self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1]) + self.disp_head = DispHead(hidden_dims[2], hidden_dim=256, output_dim=1) + + factor = 2**self.args.n_downsample + + self.mask = nn.Sequential( + nn.Conv2d(hidden_dims[2], 256, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(256, (factor**2)*9, 1, padding=0)) + + def forward(self, net, inp, corr=None, disp=None, iter08=True, iter16=True, iter32=True, update=True): + + if iter32: + net[2] = self.gru32(net[2], *(inp[2]), pool2x(net[1])) + if iter16: + if self.args.n_gru_layers > 2: + net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]), interp(net[2], net[1])) + else: + net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0])) + if iter08: + motion_features = self.encoder(disp, corr) + if self.args.n_gru_layers > 1: + net[0] = self.gru08(net[0], *(inp[0]), motion_features, interp(net[1], net[0])) + else: + net[0] = self.gru08(net[0], *(inp[0]), motion_features) + + if not update: + return net + + delta_disp = self.disp_head(net[0]) + + # scale mask to balence gradients + mask = .25 * self.mask(net[0]) + return net, mask, delta_disp + + +class ScaleBasicMultiUpdateBlock(nn.Module): + def __init__(self, args, hidden_dims=[128, 128, 128]): + super().__init__() + self.args = args + encoder_output_dim = 128 + cor_planes = len(args.scale_list) * (2*args.scale_corr_radius + 1) + self.encoder = BasicMotionEncoder(cor_planes, out_planes=encoder_output_dim) + + self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.n_gru_layers > 1)) + self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.n_gru_layers == 3) + hidden_dims[2]) + self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1]) + self.disp_head = DispHead(hidden_dims[2], hidden_dim=256, output_dim=1) + + factor = 2**self.args.n_downsample + + self.mask = nn.Sequential( + nn.Conv2d(hidden_dims[2], 256, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(256, (factor**2)*9, 1, padding=0)) + + def forward(self, net, inp, corr=None, disp=None, iter08=True, iter16=True, iter32=True, update=True): + + if iter32: + net[2] = self.gru32(net[2], *(inp[2]), pool2x(net[1])) + if iter16: + if self.args.n_gru_layers > 2: + net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]), interp(net[2], net[1])) + else: + net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0])) + if iter08: + motion_features = self.encoder(disp, corr) + if self.args.n_gru_layers > 1: + net[0] = self.gru08(net[0], *(inp[0]), motion_features, interp(net[1], net[0])) + else: + net[0] = self.gru08(net[0], *(inp[0]), motion_features) + + if not update: + return net + + x_disp = self.disp_head(net[0]) + scale_disp = F.relu6(torch.exp(.25*x_disp)) + + # scale mask to balence gradients + mask = .25 * self.mask(net[0]) + return net, mask, scale_disp diff --git a/cv/defom-stereo/core/utils/__init__.py b/cv/defom-stereo/core/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cv/defom-stereo/core/utils/augmentor.py b/cv/defom-stereo/core/utils/augmentor.py new file mode 100644 index 0000000..d1a3316 --- /dev/null +++ b/cv/defom-stereo/core/utils/augmentor.py @@ -0,0 +1,307 @@ +import numpy as np +import random +import warnings +import os +import time +from glob import glob +from skimage import color, io +from PIL import Image + +import cv2 +cv2.setNumThreads(0) +cv2.ocl.setUseOpenCL(False) + +import torch +from torchvision.transforms import ColorJitter, functional, Compose +import torch.nn.functional as F + + +def get_middlebury_images(): + root = "../datasets/Middlebury/MiddEval3" + with open(os.path.join(root, "official_train.txt"), 'r') as f: + lines = f.read().splitlines() + return sorted([os.path.join(root, 'trainingQ', f'{name}/im0.png') for name in lines]) + + +def get_eth3d_images(): + return sorted(glob('../datasets/ETH3D/two_view_training/*/im0.png')) + + +def get_kitti_images(): + return sorted(glob('..datasets/KITTI/training/image_2/*_10.png')) + + +def transfer_color(image, style_mean, style_stddev): + reference_image_lab = color.rgb2lab(image) + reference_stddev = np.std(reference_image_lab, axis=(0, 1), keepdims=True)# + 1 + reference_mean = np.mean(reference_image_lab, axis=(0, 1), keepdims=True) + + reference_image_lab = reference_image_lab - reference_mean + lamb = style_stddev/reference_stddev + style_image_lab = lamb * reference_image_lab + output_image_lab = style_image_lab + style_mean + l, a, b = np.split(output_image_lab, 3, axis=2) + l = l.clip(0, 100) + output_image_lab = np.concatenate((l, a, b), axis=2) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=UserWarning) + output_image_rgb = color.lab2rgb(output_image_lab) * 255 + return output_image_rgb + + +class AdjustGamma(object): + + def __init__(self, gamma_min, gamma_max, gain_min=1.0, gain_max=1.0): + self.gamma_min, self.gamma_max, self.gain_min, self.gain_max = gamma_min, gamma_max, gain_min, gain_max + + def __call__(self, sample): + gain = random.uniform(self.gain_min, self.gain_max) + gamma = random.uniform(self.gamma_min, self.gamma_max) + return functional.adjust_gamma(sample, gamma, gain) + + def __repr__(self): + return f"Adjust Gamma {self.gamma_min}, ({self.gamma_max}) and Gain ({self.gain_min}, {self.gain_max})" + + +class DispAugmentor: + def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True, yjitter=False, + saturation_range=[0.6, 1.4], gamma=[1, 1, 1, 1]): + + # spatial augmentation params + self.crop_size = crop_size + self.min_scale = min_scale + self.max_scale = max_scale + self.spatial_aug_prob = 1.0 + self.stretch_prob = 0.8 + self.max_stretch = 0.2 + + # flip augmentation params + self.yjitter = yjitter + self.do_flip = do_flip + self.v_flip_prob = 0.1 + + # photometric augmentation params + self.photo_aug = Compose([ColorJitter(brightness=0.4, contrast=0.4, saturation=saturation_range, hue=0.5/3.14), AdjustGamma(*gamma)]) + self.asymmetric_color_aug_prob = 0.2 + self.eraser_aug_prob = 0.5 + + def color_transform(self, img1, img2): + """ Photometric augmentation """ + + # asymmetric + if np.random.rand() < self.asymmetric_color_aug_prob: + img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8) + img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8) + + # symmetric + else: + image_stack = np.concatenate([img1, img2], axis=0) + image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) + img1, img2 = np.split(image_stack, 2, axis=0) + + return img1, img2 + + def eraser_transform(self, img1, img2, bounds=[50, 100]): + """ Occlusion augmentation """ + + ht, wd = img1.shape[:2] + if np.random.rand() < self.eraser_aug_prob: + mean_color = np.mean(img2.reshape(-1, 3), axis=0) + for _ in range(np.random.randint(1, 3)): + x0 = np.random.randint(0, wd) + y0 = np.random.randint(0, ht) + dx = np.random.randint(bounds[0], bounds[1]) + dy = np.random.randint(bounds[0], bounds[1]) + img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color + + return img1, img2 + + def spatial_transform(self, img1, img2, disp): + # randomly sample scale + ht, wd = img1.shape[:2] + min_scale = np.maximum( + (self.crop_size[0] + 8) / float(ht), + (self.crop_size[1] + 8) / float(wd)) + + scale = 2 ** np.random.uniform(self.min_scale, self.max_scale) + if scale>min_scale: + scale = np.random.uniform(min_scale, scale) + scale_x = scale + scale_y = scale + if np.random.rand() < self.stretch_prob: + scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch) + scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch) + + scale_x = np.clip(scale_x, min_scale, 2*min_scale) + scale_y = np.clip(scale_y, min_scale, 2*min_scale) + + if np.random.rand() < self.spatial_aug_prob or min_scale >= 1.0: + # rescale the images + img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) + img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) + disp = cv2.resize(disp, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) + disp = disp * scale_x + + if self.do_flip: + if np.random.rand() < self.v_flip_prob and self.do_flip == 'v': # v-flip + img1 = img1[::-1, :] + img2 = img2[::-1, :] + disp = disp[::-1, :] + + if self.yjitter: + y0 = np.random.randint(2, img1.shape[0] - self.crop_size[0] - 2) + x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1] - 0) + + y1 = y0 + np.random.randint(-2, 2 + 1) + y1 = np.clip(y1, 0, img1.shape[0] - self.crop_size[0]) + img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + img2 = img2[y1:y1 + self.crop_size[0], x0:x0 + self.crop_size[1]] + disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + + else: + y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0]) + x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1]) + + img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + + return img1, img2, disp + + def __call__(self, img1, img2, disp): + img1, img2 = self.color_transform(img1, img2) + img1, img2 = self.eraser_transform(img1, img2) + img1, img2, disp = self.spatial_transform(img1, img2, disp) + + img1 = np.ascontiguousarray(img1) + img2 = np.ascontiguousarray(img2) + disp = np.ascontiguousarray(disp) + + return img1, img2, disp + + +class SparseDispAugmentor: + def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False, yjitter=False, + saturation_range=[0.7, 1.3], gamma=[1, 1, 1, 1]): + # spatial augmentation params + self.crop_size = crop_size + self.min_scale = min_scale + self.max_scale = max_scale + self.spatial_aug_prob = 0.8 + self.stretch_prob = 0.8 + self.max_stretch = 0.2 + + # flip augmentation params + self.do_flip = do_flip + self.v_flip_prob = 0.1 + + # photometric augmentation params + self.photo_aug = Compose( + [ColorJitter(brightness=0.3, contrast=0.3, saturation=saturation_range, hue=0.3/3.14), + AdjustGamma(*gamma)]) + self.asymmetric_color_aug_prob = 0.2 + self.eraser_aug_prob = 0.5 + + def color_transform(self, img1, img2): + image_stack = np.concatenate([img1, img2], axis=0) + image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) + img1, img2 = np.split(image_stack, 2, axis=0) + return img1, img2 + + def eraser_transform(self, img1, img2): + ht, wd = img1.shape[:2] + if np.random.rand() < self.eraser_aug_prob: + mean_color = np.mean(img2.reshape(-1, 3), axis=0) + for _ in range(np.random.randint(1, 3)): + x0 = np.random.randint(0, wd) + y0 = np.random.randint(0, ht) + dx = np.random.randint(50, 100) + dy = np.random.randint(50, 100) + img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color + + return img1, img2 + + def resize_sparse_flow_map(self, disp, valid, fx=1.0, fy=1.0): + ht, wd = disp.shape[:2] + coords = np.meshgrid(np.arange(wd), np.arange(ht)) + coords = np.stack(coords, axis=-1) + + coords = coords.reshape(-1, 2).astype(np.float32) + disp = disp.reshape(-1).astype(np.float32) + valid = valid.reshape(-1).astype(np.float32) + + coords0 = coords[valid >= 1] + disp0 = disp[valid >= 1] + + ht1 = int(round(ht * fy)) + wd1 = int(round(wd * fx)) + + coords1 = coords0 * [fx, fy] + disp1 = disp0 * fx + + xx = np.round(coords1[:, 0]).astype(np.int32) + yy = np.round(coords1[:, 1]).astype(np.int32) + + v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1) + xx = xx[v] + yy = yy[v] + disp1 = disp1[v] + + disp_img = np.zeros([ht1, wd1], dtype=np.float32) + valid_img = np.zeros([ht1, wd1], dtype=np.int32) + + disp_img[yy, xx] = disp1 + valid_img[yy, xx] = 1 + + return disp_img, valid_img + + def spatial_transform(self, img1, img2, disp, valid): + # randomly sample scale + + ht, wd = img1.shape[:2] + min_scale = np.maximum( + (self.crop_size[0] + 1) / float(ht), + (self.crop_size[1] + 1) / float(wd)) + + scale = 2 ** np.random.uniform(self.min_scale, self.max_scale) + if scale>min_scale: + scale = np.random.uniform(min_scale, 2*min_scale) + scale_x = scale + scale_y = scale + + scale_x = np.clip(scale_x, min_scale, 2*min_scale) + scale_y = np.clip(scale_y, min_scale, 2*min_scale) + + if np.random.rand() < self.spatial_aug_prob or min_scale >= 1.0: + # rescale the images + img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) + img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR) + disp, valid = self.resize_sparse_flow_map(disp, valid, fx=scale_x, fy=scale_y) + + if self.do_flip: + if np.random.rand() < self.v_flip_prob and self.do_flip == 'v': # v-flip + img1 = img1[::-1, :] + img2 = img2[::-1, :] + disp = disp[::-1, :] + valid = valid[::-1, :] + + y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0]) + x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1]) + + img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + valid = valid[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + return img1, img2, disp, valid + + def __call__(self, img1, img2, disp, valid): + img1, img2 = self.color_transform(img1, img2) + img1, img2 = self.eraser_transform(img1, img2) + img1, img2, disp, valid = self.spatial_transform(img1, img2, disp, valid) + + img1 = np.ascontiguousarray(img1) + img2 = np.ascontiguousarray(img2) + disp = np.ascontiguousarray(disp) + valid = np.ascontiguousarray(valid) + + return img1, img2, disp, valid diff --git a/cv/defom-stereo/core/utils/dist_utils.py b/cv/defom-stereo/core/utils/dist_utils.py new file mode 100644 index 0000000..fdb08c7 --- /dev/null +++ b/cv/defom-stereo/core/utils/dist_utils.py @@ -0,0 +1,105 @@ +# Copyright (c) OpenMMLab. All rights reserved. +# https://github.com/open-mmlab/mmcv/blob/7540cf73ac7e5d1e14d0ffbd9b6759e83929ecfc/mmcv/runner/dist_utils.py + +import os +import subprocess + +import torch +import torch.multiprocessing as mp +from torch import distributed as dist + + +def init_dist(launcher, backend='nccl', **kwargs): + if mp.get_start_method(allow_none=True) is None: + mp.set_start_method('spawn') + if launcher == 'pytorch': + _init_dist_pytorch(backend, **kwargs) + elif launcher == 'mpi': + _init_dist_mpi(backend, **kwargs) + elif launcher == 'slurm': + _init_dist_slurm(backend, **kwargs) + else: + raise ValueError(f'Invalid launcher type: {launcher}') + + +def _init_dist_pytorch(backend, **kwargs): + # TODO: use local_rank instead of rank % num_gpus + rank = int(os.environ['RANK']) + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(rank % num_gpus) + dist.init_process_group(backend=backend, **kwargs) + + +def _init_dist_mpi(backend, **kwargs): + # TODO: use local_rank instead of rank % num_gpus + rank = int(os.environ['OMPI_COMM_WORLD_RANK']) + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(rank % num_gpus) + dist.init_process_group(backend=backend, **kwargs) + + +def _init_dist_slurm(backend, port=None): + """Initialize slurm distributed training environment. + If argument ``port`` is not specified, then the master port will be system + environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system + environment variable, then a default port ``29500`` will be used. + Args: + backend (str): Backend of torch.distributed. + port (int, optional): Master port. Defaults to None. + """ + proc_id = int(os.environ['SLURM_PROCID']) + ntasks = int(os.environ['SLURM_NTASKS']) + node_list = os.environ['SLURM_NODELIST'] + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(proc_id % num_gpus) + addr = subprocess.getoutput( + f'scontrol show hostname {node_list} | head -n1') + # specify master port + if port is not None: + os.environ['MASTER_PORT'] = str(port) + elif 'MASTER_PORT' in os.environ: + pass # use MASTER_PORT in the environment variable + else: + # 29500 is torch.distributed default port + os.environ['MASTER_PORT'] = '29500' + # use MASTER_ADDR in the environment variable if it already exists + if 'MASTER_ADDR' not in os.environ: + os.environ['MASTER_ADDR'] = addr + os.environ['WORLD_SIZE'] = str(ntasks) + os.environ['LOCAL_RANK'] = str(proc_id % num_gpus) + os.environ['RANK'] = str(proc_id) + dist.init_process_group(backend=backend) + + +def get_dist_info(): + # if (TORCH_VERSION != 'parrots' + # and digit_version(TORCH_VERSION) < digit_version('1.0')): + # initialized = dist._initialized + # else: + if dist.is_available(): + initialized = dist.is_initialized() + else: + initialized = False + if initialized: + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + return rank, world_size + + +# from DETR repo +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop('force', False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print diff --git a/cv/defom-stereo/core/utils/frame_utils.py b/cv/defom-stereo/core/utils/frame_utils.py new file mode 100644 index 0000000..17b1bbe --- /dev/null +++ b/cv/defom-stereo/core/utils/frame_utils.py @@ -0,0 +1,286 @@ +import numpy as np +from PIL import Image +from os.path import * +import re +import json +import imageio +import os +import math + +os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" +import cv2 +cv2.setNumThreads(0) +cv2.ocl.setUseOpenCL(False) + +TAG_CHAR = np.array([202021.25], np.float32) + + +def readFlow(fn): + """ Read .flo file in Middlebury format""" + # Code adapted from: + # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy + + # WARNING: this will work on little-endian architectures (eg Intel x86) only! + # print 'fn = %s'%(fn) + with open(fn, 'rb') as f: + magic = np.fromfile(f, np.float32, count=1) + if 202021.25 != magic: + print('Magic number incorrect. Invalid .flo file') + return None + else: + w = np.fromfile(f, np.int32, count=1) + h = np.fromfile(f, np.int32, count=1) + # print 'Reading %d x %d flo file\n' % (w, h) + data = np.fromfile(f, np.float32, count=2*int(w)*int(h)) + # Reshape data into 3D array (columns, rows, bands) + # The reshape here is for visualization, the original code is (w,h,2) + return np.resize(data, (int(h), int(w), 2)) + + +def readPFM(file): + file = open(file, 'rb') + + color = None + width = None + height = None + scale = None + endian = None + + header = file.readline().rstrip() + if header == b'PF': + color = True + elif header == b'Pf': + color = False + else: + raise Exception('Not a PFM file.') + + dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline()) + if dim_match: + width, height = map(int, dim_match.groups()) + else: + raise Exception('Malformed PFM header.') + + scale = float(file.readline().rstrip()) + if scale < 0: # little-endian + endian = '<' + scale = -scale + else: + endian = '>' # big-endian + + data = np.fromfile(file, endian + 'f') + shape = (height, width, 3) if color else (height, width) + + data = np.reshape(data, shape) + data = np.flipud(data) + return data + + +def writePFM(file, array): + import os + assert type(file) is str and type(array) is np.ndarray and \ + os.path.splitext(file)[1] == ".pfm" + with open(file, 'wb') as f: + H, W = array.shape + headers = ["Pf\n", f"{W} {H}\n", "-1\n"] + for header in headers: + f.write(str.encode(header)) + array = np.flip(array, axis=0).astype(np.float32) + f.write(array.tobytes()) + + +def writeFlow(filename,uv,v=None): + """ Write optical flow to file. + + If v is None, uv is assumed to contain both u and v channels, + stacked in depth. + Original code by Deqing Sun, adapted from Daniel Scharstein. + """ + nBands = 2 + + if v is None: + assert(uv.ndim == 3) + assert(uv.shape[2] == 2) + u = uv[:,:,0] + v = uv[:,:,1] + else: + u = uv + + assert(u.shape == v.shape) + height,width = u.shape + f = open(filename,'wb') + # write the header + f.write(TAG_CHAR) + np.array(width).astype(np.int32).tofile(f) + np.array(height).astype(np.int32).tofile(f) + # arrange into matrix form + tmp = np.zeros((height, width*nBands)) + tmp[:,np.arange(width)*2] = u + tmp[:,np.arange(width)*2 + 1] = v + tmp.astype(np.float32).tofile(f) + f.close() + + +def readFlowKITTI(filename): + flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR) + flow = flow[:,:,::-1].astype(np.float32) + flow, valid = flow[:, :, :2], flow[:, :, 2] + flow = (flow - 2**15) / 64.0 + return flow, valid + + +def readDispKITTI(filename): + disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0 + valid = disp > 0.0 + return disp, valid + + +def readDispInStereo2K(filename): + disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 100.0 + valid = disp > 0.0 + return disp, valid + + +def readDispVKITTI2(filename): + depth = cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH).astype(np.float32) / 100.0 + valid = depth > 0.0 + baseline = 0.532725 + focus_length = 725.0087 + disp = baseline*focus_length/(depth+1e-8) + return disp, valid + + +def readDispCreStereo(filename): + disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 32 + valid = disp > -1e-8 + return disp, valid + + +# Method taken from /n/fs/raft-depth/RAFT-Stereo/datasets/SintelStereo/sdk/python/sintel_io.py +def readDispSintelStereo(file_name): + a = np.array(Image.open(file_name)) + d_r, d_g, d_b = np.split(a, axis=2, indices_or_sections=3) + disp = (d_r.astype('float64') * 4 + d_g.astype('float64') / (2**6) + d_b.astype('float64') / (2**14))[..., 0] + mask = np.array(Image.open(file_name.replace('disparities', 'occlusions'))) + valid = ((mask == 0) & (disp > -1e-8)) + return disp, valid + + +# Method taken from https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt +def readDispFallingThings(file_name): + a = np.array(Image.open(file_name)) + with open('/'.join(file_name.split('/')[:-1] + ['_camera_settings.json']), 'r') as f: + intrinsics = json.load(f) + if 'left' in file_name: + idx = 0 + else: + idx = 1 + fx = intrinsics['camera_settings'][idx]['intrinsic_settings']['fx'] + disp = (fx * 6.0 * 100) / a.astype(np.float32) + valid = disp > -1e-8 + return disp, valid + + +# Method taken from https://github.com/castacks/tartanair_tools/blob/master/data_type.md +def readDispTartanAir(file_name): + depth = np.load(file_name) + disp = 80.0 / depth + valid = disp > -1e-8 + return disp, valid + + +def readDispBooster(file_name): + disp = np.load(file_name) + valid = disp > 0 + return disp, valid + + +def readDisp3DKenBurns(file_name): + depth = cv2.imread(file_name, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) + meta_file_name = file_name.replace('-depth', '')[:-7]+'-meta.json' + fltFov = json.loads(open(meta_file_name, 'r').read())['fltFov'] + fltFocal = 0.5 * 512 * math.tan(math.radians(90.0) - (0.5 * math.radians(fltFov))) + fltBaseline = 40.0 + disp = (fltFocal * fltBaseline) / depth + valid = disp > 0 + return disp, valid + + +def readDispMiddlebury0(file_name): + if basename(file_name) == 'disp0GT.pfm': + disp = readPFM(file_name).astype(np.float32) + assert len(disp.shape) == 2 + nocc_pix = file_name.replace('disp0GT.pfm', 'mask0nocc.png') + assert exists(nocc_pix) + nocc_pix = imageio.imread(nocc_pix) == 255 + assert np.any(nocc_pix) + return disp, nocc_pix + elif basename(file_name) == 'disp1GT.pfm': + disp = readPFM(file_name).astype(np.float32) + assert len(disp.shape) == 2 + nocc_pix = file_name.replace('disp1GT.pfm', 'mask1nocc.png') + assert exists(nocc_pix) + nocc_pix = imageio.imread(nocc_pix) == 255 + assert np.any(nocc_pix) + return disp, nocc_pix + elif basename(file_name) == 'disp0.pfm': + disp = readPFM(file_name).astype(np.float32) + valid = disp < 1e3 + return disp, valid + elif basename(file_name) == 'disp1.pfm': + disp = readPFM(file_name).astype(np.float32) + valid = disp < 1e3 + return disp, valid + elif splitext(file_name)[-1] == '.png': + disp = np.array(Image.open(file_name)).astype(np.float32) + valid = disp > 0.0 + return disp, valid + + +def readDispMiddlebury(file_name): + if basename(file_name) == 'disp0GT.pfm': + disp = readPFM(file_name).astype(np.float32) + return disp, disp<1e3 + elif basename(file_name) == 'disp1GT.pfm': + disp = readPFM(file_name).astype(np.float32) + return disp, disp<1e3 + elif basename(file_name) == 'disp0.pfm': + disp = readPFM(file_name).astype(np.float32) + valid = disp < 1e3 + return disp, valid + elif basename(file_name) == 'disp1.pfm': + disp = readPFM(file_name).astype(np.float32) + valid = disp < 1e3 + return disp, valid + elif splitext(file_name)[-1] == '.png': + disp = np.array(Image.open(file_name)).astype(np.float32) + valid = disp > 0.0 + return disp, valid + + +def writeFlowKITTI(filename, uv): + uv = 64.0 * uv + 2**15 + valid = np.ones([uv.shape[0], uv.shape[1], 1]) + uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16) + cv2.imwrite(filename, uv[..., ::-1]) + + +def read_gen(file_name, pil=False): + ext = splitext(file_name)[-1] + if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg': + return Image.open(file_name) + elif ext == '.bin' or ext == '.raw': + return np.load(file_name) + elif ext == '.flo': + return readFlow(file_name).astype(np.float32) + elif ext == '.pfm': + flow = readPFM(file_name).astype(np.float32) + if len(flow.shape) == 2: + return flow + else: + return flow[:, :, :-1] + elif ext == '.exr': + disp = cv2.imread(file_name, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) + if len(disp.shape) > 2: + disp = disp[..., 0] + return disp + return [] diff --git a/cv/defom-stereo/core/utils/utils.py b/cv/defom-stereo/core/utils/utils.py new file mode 100644 index 0000000..d691470 --- /dev/null +++ b/cv/defom-stereo/core/utils/utils.py @@ -0,0 +1,242 @@ +import torch +import torch.nn.functional as F +import numpy as np +from scipy import interpolate +import glob +import os.path as osp + + +def get_danv2_io_size(h, w, nds, max_i_size=2688, multiple_of=14): + """compute the input and output sizes of danv2 network""" + danv2_oh, danv2_ow = h//2**nds, w//2**nds + danv2_io_factor = 3.5 # more precise, 14/8=3.5 + ih, iw = danv2_io_factor*danv2_oh, danv2_io_factor*danv2_ow + ih = int(np.ceil(ih / multiple_of) * multiple_of) + iw = int(np.ceil(iw / multiple_of) * multiple_of) + + max_i_size = int(np.floor(max_i_size / multiple_of) * multiple_of) + + if ih <= max_i_size and iw <= max_i_size: + danv2_ih, danv2_iw = ih, iw + else: + factor_h = max_i_size/ih + factor_w = max_i_size/iw + + if factor_w > factor_h: + danv2_ih = max_i_size + danv2_iw = int(np.ceil(factor_h * iw / multiple_of) * multiple_of) + else: + danv2_iw = max_i_size + danv2_ih = int(np.ceil(factor_w * ih / multiple_of) * multiple_of) + + return danv2_ih, danv2_iw, danv2_oh, danv2_ow + + +class InputPadder: + """ Pads images such that dimensions are divisible by 8 """ + def __init__(self, dims, mode='sintel', divis_by=8): + self.ht, self.wd = dims[-2:] + pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by + pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by + if mode == 'sintel': + self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2] + else: + self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht] + + def pad(self, *inputs): + assert all((x.ndim == 4) for x in inputs) + return [F.pad(x, self._pad, mode='replicate') for x in inputs] + + def unpad(self, x): + assert x.ndim == 4 + ht, wd = x.shape[-2:] + c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]] + return x[..., c[0]:c[1], c[2]:c[3]] + + +def forward_interpolate(flow): + flow = flow.detach().cpu().numpy() + dx, dy = flow[0], flow[1] + + ht, wd = dx.shape + x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) + + x1 = x0 + dx + y1 = y0 + dy + + x1 = x1.reshape(-1) + y1 = y1.reshape(-1) + dx = dx.reshape(-1) + dy = dy.reshape(-1) + + valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) + x1 = x1[valid] + y1 = y1[valid] + dx = dx[valid] + dy = dy[valid] + + flow_x = interpolate.griddata( + (x1, y1), dx, (x0, y0), method='nearest', fill_value=0) + + flow_y = interpolate.griddata( + (x1, y1), dy, (x0, y0), method='nearest', fill_value=0) + + flow = np.stack([flow_x, flow_y], axis=0) + return torch.from_numpy(flow).float() + + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + xgrid, ygrid = coords.split([1, 1], dim=-1) + xgrid = 2*xgrid/(W-1) - 1 + if H > 1: + ygrid = 2*ygrid/(H-1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + img = F.grid_sample(img, grid, align_corners=True) + # img = bilinear_grid_sample(img, grid, align_corners=True) + + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + + +def coords_grid(batch, ht, wd): + coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +def upflow(flow, factor=8, mode='bilinear', sacle=True): + new_size = (factor * flow.shape[2], factor * flow.shape[3]) + if sacle: + return factor * F.interpolate(flow, size=new_size, mode=mode, align_corners=True) + else: + return F.interpolate(flow, size=new_size, mode=mode, align_corners=True) + + +def gauss_blur(input, N=5, std=1): + B, D, H, W = input.shape + x, y = torch.meshgrid(torch.arange(N).float() - N//2, torch.arange(N).float() - N//2) + unnormalized_gaussian = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * std ** 2)) + weights = unnormalized_gaussian / unnormalized_gaussian.sum().clamp(min=1e-4) + weights = weights.view(1, 1, N, N).to(input) + output = F.conv2d(input.reshape(B*D, 1, H, W), weights, padding=N//2) + return output.view(B, D, H, W) + + +# Ref: https://zenn.dev/pinto0309/scraps/7d4032067d0160 +def bilinear_grid_sample(im, grid, align_corners=False): + """Given an input and a flow-field grid, computes the output using input + values and pixel locations from grid. Supported only bilinear interpolation + method to sample the input pixels. + + Args: + im (torch.Tensor): Input feature map, shape (N, C, H, W) + grid (torch.Tensor): Point coordinates, shape (N, Hg, Wg, 2) + align_corners {bool}: If set to True, the extrema (-1 and 1) are + considered as referring to the center points of the input’s + corner pixels. If set to False, they are instead considered as + referring to the corner points of the input’s corner pixels, + making the sampling more resolution agnostic. + + Returns: + torch.Tensor: A tensor with sampled points, shape (N, C, Hg, Wg) + """ + n, c, h, w = im.shape + gn, gh, gw, _ = grid.shape + assert n == gn + + x = grid[:, :, :, 0] + y = grid[:, :, :, 1] + + if align_corners: + x = ((x + 1) / 2) * (w - 1) + y = ((y + 1) / 2) * (h - 1) + else: + x = ((x + 1) * w - 1) / 2 + y = ((y + 1) * h - 1) / 2 + + x = x.view(n, -1) + y = y.view(n, -1) + + x0 = torch.floor(x).long() + y0 = torch.floor(y).long() + x1 = x0 + 1 + y1 = y0 + 1 + + wa = ((x1 - x) * (y1 - y)).unsqueeze(1) + wb = ((x1 - x) * (y - y0)).unsqueeze(1) + wc = ((x - x0) * (y1 - y)).unsqueeze(1) + wd = ((x - x0) * (y - y0)).unsqueeze(1) + + # Apply default for grid_sample function zero padding + im_padded = torch.nn.functional.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0) + padded_h = h + 2 + padded_w = w + 2 + # save points positions after padding + x0, x1, y0, y1 = x0 + 1, x1 + 1, y0 + 1, y1 + 1 + + # Clip coordinates to padded image size + x0 = torch.where(x0 < 0, torch.tensor(0, device=im.device), x0) + x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x0) + x1 = torch.where(x1 < 0, torch.tensor(0, device=im.device), x1) + x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x1) + y0 = torch.where(y0 < 0, torch.tensor(0, device=im.device), y0) + y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y0) + y1 = torch.where(y1 < 0, torch.tensor(0, device=im.device), y1) + y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y1) + + im_padded = im_padded.view(n, c, -1) + + x0_y0 = (x0 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1) + x0_y1 = (x0 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1) + x1_y0 = (x1 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1) + x1_y1 = (x1 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1) + + Ia = torch.gather(im_padded, 2, x0_y0) + Ib = torch.gather(im_padded, 2, x0_y1) + Ic = torch.gather(im_padded, 2, x1_y0) + Id = torch.gather(im_padded, 2, x1_y1) + + return (Ia * wa + Ib * wb + Ic * wc + Id * wd).reshape(n, c, gh, gw) + + +def read_kitti_calib_file(path): + """Read KITTI calibration file + (from https://github.com/hunse/kitti) + """ + float_chars = set("0123456789.e+- ") + data = {} + with open(path, 'r') as f: + for line in f.readlines(): + key, value = line.split(':', 1) + value = value.strip() + data[key] = value + if float_chars.issuperset(value): + # try to cast to float array + try: + data[key] = np.array(list(map(float, value.split(' ')))) + except ValueError: + # casting error: data[key] already eq. value, so pass + pass + + return data + + +# from https://github.com/ozendelait/rvc_devkit/blob/master/stereo/stereo_devkit.py +def ReadMiddlebury2014CalibFile(path): + result = dict() + with open(path, 'rb') as calib_file: + for line in calib_file.readlines(): + line = line.decode('UTF-8').rstrip('\n') + if len(line) == 0: + continue + eq_pos = line.find('=') + if eq_pos < 0: + raise Exception('Cannot parse Middlebury 2014 calib file: ' + path) + result[line[:eq_pos]] = line[eq_pos + 1:] + return result diff --git a/cv/defom-stereo/demo.py b/cv/defom-stereo/demo.py new file mode 100644 index 0000000..e82a204 --- /dev/null +++ b/cv/defom-stereo/demo.py @@ -0,0 +1,89 @@ +import sys +sys.path.append('core') + +import argparse +import glob +import numpy as np +import torch +from tqdm import tqdm +from pathlib import Path +from core.defom_stereo import DEFOMStereo +from utils.utils import InputPadder +from PIL import Image +from matplotlib import pyplot as plt + + +DEVICE = 'cuda' + +def load_image(imfile): + img = np.array(Image.open(imfile)).astype(np.uint8) + img = torch.from_numpy(img).permute(2, 0, 1).float() + return img[None].to(DEVICE) + +def demo(args): + model = DEFOMStereo(args) + checkpoint = torch.load(args.restore_ckpt, map_location='cuda') + if 'model' in checkpoint: + model.load_state_dict(checkpoint['model']) + else: + model.load_state_dict(checkpoint) + + model.to(DEVICE) + model.eval() + + output_directory = Path(args.output_directory) + output_directory.mkdir(exist_ok=True) + + with torch.no_grad(): + left_images = sorted(glob.glob(args.left_imgs, recursive=True)) + right_images = sorted(glob.glob(args.right_imgs, recursive=True)) + print(f"Found {len(left_images)} images. Saving files to {output_directory}/") + + for (imfile1, imfile2) in tqdm(list(zip(left_images, right_images))): + image1 = load_image(imfile1) + image2 = load_image(imfile2) + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with torch.no_grad(): + disp_pr = model(image1, image2, iters=args.valid_iters, scale_iters=args.scale_iters, test_mode=True) + disp_pr = padder.unpad(disp_pr).cpu().squeeze().numpy() + + file_stem = imfile1.split('/')[-1].split('_')[0]+'_'+args.restore_ckpt.split('/')[-1][:-4] + if args.save_numpy: + np.save(output_directory / f"{file_stem}.npy", disp_pr) + plt.imsave(output_directory / f"{file_stem}.png", disp_pr, cmap='jet') + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--restore_ckpt', help="restore checkpoint", required=True) + parser.add_argument('--save_numpy', action='store_true', help='save output as numpy arrays') + parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default="demo/*_left.png") + parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default="demo/*_right.png") + parser.add_argument('--output_directory', help="directory to save output", default="demo") + parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') + parser.add_argument('--valid_iters', type=int, default=32, help='number of flow-field updates during forward pass') + parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.") + + # Architecture choices + parser.add_argument('--dinov2_encoder', type=str, default='vitl', choices=['vits', 'vitb', 'vitl', 'vitg']) + parser.add_argument('--idepth_scale', type=float, default=0.5, help="the scale of inverse depth to initialize disparity") + parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions") + parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation") + parser.add_argument('--shared_backbone', action='store_true', help="use a single backbone for the context and feature encoders") + parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid") + parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid") + parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0], + help='the list of scaling factors of disparity') + parser.add_argument('--scale_corr_radius', type=int, default=2, + help="width of the correlation pyramid for scaled disparity") + + parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], help="resolution of the disparity field (1/2^K)") + parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder") + parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels") + + args = parser.parse_args() + + demo(args) diff --git a/cv/defom-stereo/demo/1_left.png b/cv/defom-stereo/demo/1_left.png new file mode 100644 index 0000000..94c7b6b Binary files /dev/null and b/cv/defom-stereo/demo/1_left.png differ diff --git a/cv/defom-stereo/demo/1_right.png b/cv/defom-stereo/demo/1_right.png new file mode 100644 index 0000000..0414183 Binary files /dev/null and b/cv/defom-stereo/demo/1_right.png differ diff --git a/cv/defom-stereo/demo/2_left.png b/cv/defom-stereo/demo/2_left.png new file mode 100644 index 0000000..f622ca8 Binary files /dev/null and b/cv/defom-stereo/demo/2_left.png differ diff --git a/cv/defom-stereo/demo/2_right.png b/cv/defom-stereo/demo/2_right.png new file mode 100644 index 0000000..6024f69 Binary files /dev/null and b/cv/defom-stereo/demo/2_right.png differ diff --git a/cv/defom-stereo/demo/3_left.png b/cv/defom-stereo/demo/3_left.png new file mode 100644 index 0000000..4a7b238 Binary files /dev/null and b/cv/defom-stereo/demo/3_left.png differ diff --git a/cv/defom-stereo/demo/3_right.png b/cv/defom-stereo/demo/3_right.png new file mode 100644 index 0000000..81803ad Binary files /dev/null and b/cv/defom-stereo/demo/3_right.png differ diff --git a/cv/defom-stereo/demo/Thumbs.db b/cv/defom-stereo/demo/Thumbs.db new file mode 100644 index 0000000..3221794 Binary files /dev/null and b/cv/defom-stereo/demo/Thumbs.db differ diff --git a/cv/defom-stereo/depth_anything_v2/__init__.py b/cv/defom-stereo/depth_anything_v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cv/defom-stereo/depth_anything_v2/dinov2.py b/cv/defom-stereo/depth_anything_v2/dinov2.py new file mode 100644 index 0000000..83d2508 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2.py @@ -0,0 +1,415 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +from functools import partial +import math +import logging +from typing import Sequence, Tuple, Union, Callable + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from torch.nn.init import trunc_normal_ + +from .dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block + + +logger = logging.getLogger("dinov2") + + +def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=None, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + drop_path_rate (float): stochastic depth rate + drop_path_uniform (bool): apply uniform drop rate across blocks + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings + """ + super().__init__() + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + + self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) + num_patches = self.patch_embed.num_patches + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + if ffn_layer == "mlp": + logger.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + logger.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + logger.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + ) + for i in range(depth) + ] + if block_chunks > 0: + self.chunked_blocks = True + chunked_blocks = [] + chunksize = depth // block_chunks + for i in range(0, depth, chunksize): + # this is to keep the block index consistent if we chunk the block list + chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize]) + self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks]) + else: + self.chunked_blocks = False + self.blocks = nn.ModuleList(blocks_list) + + self.norm = norm_layer(embed_dim) + self.head = nn.Identity() + + self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) + + self.init_weights() + + def init_weights(self): + trunc_normal_(self.pos_embed, std=0.02) + nn.init.normal_(self.cls_token, std=1e-6) + if self.register_tokens is not None: + nn.init.normal_(self.register_tokens, std=1e-6) + named_apply(init_weights_vit_timm, self) + + def interpolate_pos_encoding(self, x, w, h): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + # we add a small number to avoid floating point error in the interpolation + # see discussion at https://github.com/facebookresearch/dino/issues/8 + # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0 + w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset + # w0, h0 = w0 + 0.1, h0 + 0.1 + + sqrt_N = math.sqrt(N) + sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2), + scale_factor=(sx, sy), + # (int(w0), int(h0)), # to solve the upsampling shape issue + mode="bicubic", + antialias=self.interpolate_antialias + ) + + assert int(w0) == patch_pos_embed.shape[-2] + assert int(h0) == patch_pos_embed.shape[-1] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) + + def prepare_tokens_with_masks(self, x, masks=None): + B, nc, w, h = x.shape + x = self.patch_embed(x) + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + + x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) + x = x + self.interpolate_pos_encoding(x, w, h) + + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + + return x + + def forward_features_list(self, x_list, masks_list): + x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)] + for blk in self.blocks: + x = blk(x) + + all_x = x + output = [] + for x, masks in zip(all_x, masks_list): + x_norm = self.norm(x) + output.append( + { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + ) + return output + + def forward_features(self, x, masks=None): + if isinstance(x, list): + return self.forward_features_list(x, masks) + + x = self.prepare_tokens_with_masks(x, masks) + + for blk in self.blocks: + x = blk(x) + + x_norm = self.norm(x) + return { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + + def _get_intermediate_layers_not_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + # If n is an int, take the n last blocks. If it's a list, take them + output, total_block_len = [], len(self.blocks) + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for i, blk in enumerate(self.blocks): + x = blk(x) + if i in blocks_to_take: + output.append(x) + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def _get_intermediate_layers_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + output, i, total_block_len = [], 0, len(self.blocks[-1]) + # If n is an int, take the n last blocks. If it's a list, take them + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for block_chunk in self.blocks: + for blk in block_chunk[i:]: # Passing the nn.Identity() + x = blk(x) + if i in blocks_to_take: + output.append(x) + i += 1 + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + norm=True + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + if self.chunked_blocks: + outputs = self._get_intermediate_layers_chunked(x, n) + else: + outputs = self._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs = [self.norm(out) for out in outputs] + class_tokens = [out[:, 0] for out in outputs] + outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs] + if reshape: + B, _, w, h = x.shape + outputs = [ + out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous() + for out in outputs + ] + if return_class_token: + return tuple(zip(outputs, class_tokens)) + return tuple(outputs) + + def forward(self, *args, is_training=False, **kwargs): + ret = self.forward_features(*args, **kwargs) + if is_training: + return ret + else: + return self.head(ret["x_norm_clstoken"]) + + +def init_weights_vit_timm(module: nn.Module, name: str = ""): + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + + +def vit_small(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs): + """ + Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 + """ + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=40, + num_heads=24, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def DINOv2(model_name): + model_zoo = { + "vits": vit_small, + "vitb": vit_base, + "vitl": vit_large, + "vitg": vit_giant2 + } + + return model_zoo[model_name]( + img_size=518, + patch_size=14, + init_values=1.0, + ffn_layer="mlp" if model_name != "vitg" else "swiglufused", + block_chunks=0, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1 + ) diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/__init__.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/__init__.py new file mode 100644 index 0000000..8120f4b --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from .mlp import Mlp +from .patch_embed import PatchEmbed +from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused +from .block import NestedTensorBlock +from .attention import MemEffAttention diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/attention.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/attention.py new file mode 100644 index 0000000..ab91702 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/attention.py @@ -0,0 +1,82 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import logging + +from torch import Tensor +from torch import nn + + +logger = logging.getLogger("dinov2") + + +try: + from xformers.ops import memory_efficient_attention, unbind, fmha + + XFORMERS_AVAILABLE = True +except ImportError: + logger.warning("xFormers not available") + XFORMERS_AVAILABLE = False + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: Tensor) -> Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + + q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] + attn = q @ k.transpose(-2, -1) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class MemEffAttention(Attention): + def forward(self, x: Tensor, attn_bias=None) -> Tensor: + if not XFORMERS_AVAILABLE: + assert attn_bias is None, "xFormers is required for nested tensors usage" + return super().forward(x) + + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + + q, k, v = unbind(qkv, 2) + + x = memory_efficient_attention(q, k, v, attn_bias=attn_bias) + x = x.reshape([B, N, C]) + + x = self.proj(x) + x = self.proj_drop(x) + return x + diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/block.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/block.py new file mode 100644 index 0000000..25488f5 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/block.py @@ -0,0 +1,252 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +import logging +from typing import Callable, List, Any, Tuple, Dict + +import torch +from torch import nn, Tensor + +from .attention import Attention, MemEffAttention +from .drop_path import DropPath +from .layer_scale import LayerScale +from .mlp import Mlp + + +logger = logging.getLogger("dinov2") + + +try: + from xformers.ops import fmha + from xformers.ops import scaled_index_add, index_select_cat + + XFORMERS_AVAILABLE = True +except ImportError: + logger.warning("xFormers not available") + XFORMERS_AVAILABLE = False + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + ) -> None: + super().__init__() + # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") + self.norm1 = norm_layer(dim) + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + ) + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.sample_drop_ratio = drop_path + + def forward(self, x: Tensor) -> Tensor: + def attn_residual_func(x: Tensor) -> Tensor: + return self.ls1(self.attn(self.norm1(x))) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + if self.training and self.sample_drop_ratio > 0.1: + # the overhead is compensated only for a drop path rate larger than 0.1 + x = drop_add_residual_stochastic_depth( + x, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + x = drop_add_residual_stochastic_depth( + x, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + elif self.training and self.sample_drop_ratio > 0.0: + x = x + self.drop_path1(attn_residual_func(x)) + x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 + else: + x = x + attn_residual_func(x) + x = x + ffn_residual_func(x) + return x + + +def drop_add_residual_stochastic_depth( + x: Tensor, + residual_func: Callable[[Tensor], Tensor], + sample_drop_ratio: float = 0.0, +) -> Tensor: + # 1) extract subset using permutation + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + x_subset = x[brange] + + # 2) apply residual_func to get residual + residual = residual_func(x_subset) + + x_flat = x.flatten(1) + residual = residual.flatten(1) + + residual_scale_factor = b / sample_subset_size + + # 3) add the residual + x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor) + return x_plus_residual.view_as(x) + + +def get_branges_scales(x, sample_drop_ratio=0.0): + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + residual_scale_factor = b / sample_subset_size + return brange, residual_scale_factor + + +def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None): + if scaling_vector is None: + x_flat = x.flatten(1) + residual = residual.flatten(1) + x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor) + else: + x_plus_residual = scaled_index_add( + x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor + ) + return x_plus_residual + + +attn_bias_cache: Dict[Tuple, Any] = {} + + +def get_attn_bias_and_cat(x_list, branges=None): + """ + this will perform the index select, cat the tensors, and provide the attn_bias from cache + """ + batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list] + all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list)) + if all_shapes not in attn_bias_cache.keys(): + seqlens = [] + for b, x in zip(batch_sizes, x_list): + for _ in range(b): + seqlens.append(x.shape[1]) + attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens) + attn_bias._batch_sizes = batch_sizes + attn_bias_cache[all_shapes] = attn_bias + + if branges is not None: + cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1]) + else: + tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list) + cat_tensors = torch.cat(tensors_bs1, dim=1) + + return attn_bias_cache[all_shapes], cat_tensors + + +def drop_add_residual_stochastic_depth_list( + x_list: List[Tensor], + residual_func: Callable[[Tensor, Any], Tensor], + sample_drop_ratio: float = 0.0, + scaling_vector=None, +) -> Tensor: + # 1) generate random set of indices for dropping samples in the batch + branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list] + branges = [s[0] for s in branges_scales] + residual_scale_factors = [s[1] for s in branges_scales] + + # 2) get attention bias and index+concat the tensors + attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges) + + # 3) apply residual_func to get residual, and split the result + residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore + + outputs = [] + for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors): + outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x)) + return outputs + + +class NestedTensorBlock(Block): + def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]: + """ + x_list contains a list of tensors to nest together and run + """ + assert isinstance(self.attn, MemEffAttention) + + if self.training and self.sample_drop_ratio > 0.0: + + def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.attn(self.norm1(x), attn_bias=attn_bias) + + def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.mlp(self.norm2(x)) + + x_list = drop_add_residual_stochastic_depth_list( + x_list, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None, + ) + x_list = drop_add_residual_stochastic_depth_list( + x_list, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None, + ) + return x_list + else: + + def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias)) + + def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + attn_bias, x = get_attn_bias_and_cat(x_list) + x = x + attn_residual_func(x, attn_bias=attn_bias) + x = x + ffn_residual_func(x) + return attn_bias.split(x) + + def forward(self, x_or_x_list): + if isinstance(x_or_x_list, Tensor): + return super().forward(x_or_x_list) + elif isinstance(x_or_x_list, list): + assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage" + return self.forward_nested(x_or_x_list) + else: + raise AssertionError diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/drop_path.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/drop_path.py new file mode 100644 index 0000000..af05625 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/drop_path.py @@ -0,0 +1,35 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py + + +from torch import nn + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0: + random_tensor.div_(keep_prob) + output = x * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/layer_scale.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/layer_scale.py new file mode 100644 index 0000000..ca5daa5 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/layer_scale.py @@ -0,0 +1,28 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 + +from typing import Union + +import torch +from torch import Tensor +from torch import nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/mlp.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/mlp.py new file mode 100644 index 0000000..5e4b315 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/mlp.py @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py + + +from typing import Callable, Optional + +from torch import Tensor, nn + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/patch_embed.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/patch_embed.py new file mode 100644 index 0000000..574abe4 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/patch_embed.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, Optional, Tuple, Union + +from torch import Tensor +import torch.nn as nn + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" + assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops diff --git a/cv/defom-stereo/depth_anything_v2/dinov2_layers/swiglu_ffn.py b/cv/defom-stereo/depth_anything_v2/dinov2_layers/swiglu_ffn.py new file mode 100644 index 0000000..b3324b2 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dinov2_layers/swiglu_ffn.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Callable, Optional + +from torch import Tensor, nn +import torch.nn.functional as F + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +try: + from xformers.ops import SwiGLU + + XFORMERS_AVAILABLE = True +except ImportError: + SwiGLU = SwiGLUFFN + XFORMERS_AVAILABLE = False + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/cv/defom-stereo/depth_anything_v2/dpt.py b/cv/defom-stereo/depth_anything_v2/dpt.py new file mode 100644 index 0000000..c03147b --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/dpt.py @@ -0,0 +1,309 @@ +import cv2 +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.transforms import Compose + +from .dinov2 import DINOv2 +from .util.blocks import FeatureFusionBlock, _make_scratch +from .util.transform import Resize, NormalizeImage, PrepareForNet + + +def _make_fusion_block(features, use_bn, size=None): + return FeatureFusionBlock( + features, + nn.ReLU(False), + deconv=False, + bn=use_bn, + expand=False, + align_corners=True, + size=size, + ) + + +class ConvBlock(nn.Module): + def __init__(self, in_feature, out_feature): + super().__init__() + + self.conv_block = nn.Sequential( + nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(out_feature), + nn.ReLU(True) + ) + + def forward(self, x): + return self.conv_block(x) + + +class DPTHead(nn.Module): + def __init__( + self, + in_channels, + features=256, + use_bn=False, + out_channels=[256, 512, 1024, 1024], + use_clstoken=False, + ): + super(DPTHead, self).__init__() + + self.use_clstoken = use_clstoken + + self.projects = nn.ModuleList([ + nn.Conv2d( + in_channels=in_channels, + out_channels=out_channel, + kernel_size=1, + stride=1, + padding=0, + ) for out_channel in out_channels + ]) + + self.resize_layers = nn.ModuleList([ + nn.ConvTranspose2d( + in_channels=out_channels[0], + out_channels=out_channels[0], + kernel_size=4, + stride=4, + padding=0), + nn.ConvTranspose2d( + in_channels=out_channels[1], + out_channels=out_channels[1], + kernel_size=2, + stride=2, + padding=0), + nn.Identity(), + nn.Conv2d( + in_channels=out_channels[3], + out_channels=out_channels[3], + kernel_size=3, + stride=2, + padding=1) + ]) + + if use_clstoken: + self.readout_projects = nn.ModuleList() + for _ in range(len(self.projects)): + self.readout_projects.append( + nn.Sequential( + nn.Linear(2 * in_channels, in_channels), + nn.GELU())) + + self.scratch = _make_scratch( + out_channels, + features, + groups=1, + expand=False, + ) + + self.scratch.stem_transpose = None + + self.scratch.refinenet1 = _make_fusion_block(features, use_bn) + self.scratch.refinenet2 = _make_fusion_block(features, use_bn) + self.scratch.refinenet3 = _make_fusion_block(features, use_bn) + self.scratch.refinenet4 = _make_fusion_block(features, use_bn) + + head_features_1 = features + head_features_2 = 32 + + self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1) + + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1), + nn.ReLU(True), + nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), + nn.ReLU(True), + nn.Identity(), + ) + + def forward(self, out_features, patch_h, patch_w, out_h, out_w): + bs = out_features[0][0].shape[0] + out = [] + for i, x in enumerate(out_features): + if self.use_clstoken: + x, cls_token = x[0][:bs//2], x[1][:bs//2] + readout = cls_token.unsqueeze(1).expand_as(x) + x = self.readout_projects[i](torch.cat((x, readout), -1)) + else: + x = x[0][:bs//2] + + x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w)) + + x = self.projects[i](x) + x = self.resize_layers[i](x) + + out.append(x) + + layer_1, layer_2, layer_3, layer_4 = out + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:]) + path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:]) + path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:]) + path_1 = self.scratch.refinenet1(path_2, layer_1_rn) + + out = self.scratch.output_conv1(path_1) + out = F.interpolate(out, (out_h, out_w), mode="bilinear", align_corners=True) + + idepth = self.scratch.output_conv2(out) + return idepth + + +class DPTFeat(nn.Module): + def __init__( + self, + in_channels, + features=256, + use_bn=False, + out_channels=[256, 512, 1024, 1024], + use_clstoken=False, + ): + super(DPTFeat, self).__init__() + + self.use_clstoken = use_clstoken + + self.projects = nn.ModuleList([ + nn.Conv2d( + in_channels=in_channels, + out_channels=out_channel, + kernel_size=1, + stride=1, + padding=0, + ) for out_channel in out_channels + ]) + + self.resize_layers = nn.ModuleList([ + nn.ConvTranspose2d( + in_channels=out_channels[0], + out_channels=out_channels[0], + kernel_size=4, + stride=4, + padding=0), + nn.ConvTranspose2d( + in_channels=out_channels[1], + out_channels=out_channels[1], + kernel_size=2, + stride=2, + padding=0), + nn.Identity(), + nn.Conv2d( + in_channels=out_channels[3], + out_channels=out_channels[3], + kernel_size=3, + stride=2, + padding=1) + ]) + + if use_clstoken: + self.readout_projects = nn.ModuleList() + for _ in range(len(self.projects)): + self.readout_projects.append( + nn.Sequential( + nn.Linear(2 * in_channels, in_channels), + nn.GELU())) + + self.scratch = _make_scratch( + out_channels, + features, + groups=1, + expand=False, + ) + + self.scratch.stem_transpose = None + + self.scratch.refinenet1 = _make_fusion_block(features, use_bn) + self.scratch.refinenet2 = _make_fusion_block(features, use_bn) + self.scratch.refinenet3 = _make_fusion_block(features, use_bn) + self.scratch.refinenet4 = _make_fusion_block(features, use_bn) + + def forward(self, out_features, patch_h, patch_w, out_h, out_w): + bs = out_features[0][0].shape[0] + out = [] + for i, x in enumerate(out_features): + if self.use_clstoken: + x, cls_token = x[0], x[1] + readout = cls_token.unsqueeze(1).expand_as(x) + x = self.readout_projects[i](torch.cat((x, readout), -1)) + else: + x = x[0] + + x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w)) + + x = self.projects[i](x) + x = self.resize_layers[i](x) + + out.append(x) + + layer_1, layer_2, layer_3, layer_4 = out + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + layer_1_rn = F.interpolate(layer_1_rn, (out_h, out_w), mode="bilinear", align_corners=True) + layer_2_rn = F.interpolate(layer_2_rn, (out_h // 2, out_w // 2), mode="bilinear", align_corners=True) + layer_3_rn = F.interpolate(layer_3_rn, (out_h // 4, out_w // 4), mode="bilinear", align_corners=True) + layer_4_rn = F.interpolate(layer_4_rn, (out_h//8, out_w//8), mode="bilinear", align_corners=True) + + out_features = [layer_1_rn[:bs//2], layer_2_rn[:bs//2], layer_3_rn[:bs//2]] + + path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:]) + path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:]) + path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:]) + path_1 = self.scratch.refinenet1(path_2, layer_1_rn) + + return out_features, path_1[:bs//2], path_1[bs//2:] + + +class DepthAnythingV2(nn.Module): + def __init__( + self, + encoder='vitl', + features=256, + out_channels=[256, 512, 1024, 1024], + use_bn=False, + use_clstoken=False, + ): + super(DepthAnythingV2, self).__init__() + + self.intermediate_layer_idx = { + 'vits': [2, 5, 8, 11], + 'vitb': [2, 5, 8, 11], + 'vitl': [4, 11, 17, 23], + 'vitg': [9, 19, 29, 39] + } + self.encoder = encoder + self.pretrained = DINOv2(model_name=encoder) + + self.depth_head = DPTHead(self.pretrained.embed_dim, features, use_bn, + out_channels=out_channels, use_clstoken=use_clstoken) + self.depth_feat = DPTFeat(self.pretrained.embed_dim, features, use_bn, + out_channels=out_channels, use_clstoken=use_clstoken) + + + def forward(self, x, out_h, out_w): + patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14 + + features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True) + + d_features, left_feat, right_feat = self.depth_feat(features, patch_h, patch_w, out_h, out_w) + idepth = self.depth_head(features, patch_h, patch_w, out_h, out_w) + + return d_features, left_feat, right_feat, idepth + + + @torch.no_grad() + def forward_test(self, x, out_h, out_w): + patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14 + + features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], + return_class_token=True) + + d_features, left_feat, right_feat = self.depth_feat(features, patch_h, patch_w, out_h, out_w) + idepth = self.depth_head(features, patch_h, patch_w, out_h, out_w) + + return d_features, left_feat, right_feat, idepth diff --git a/cv/defom-stereo/depth_anything_v2/util/blocks.py b/cv/defom-stereo/depth_anything_v2/util/blocks.py new file mode 100644 index 0000000..d033067 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/util/blocks.py @@ -0,0 +1,149 @@ +import torch.nn as nn + + +def _make_scratch(in_shape, out_shape, groups=1, expand=False): + scratch = nn.Module() + + out_shape1 = out_shape + out_shape2 = out_shape + out_shape3 = out_shape + if len(in_shape) >= 4: + out_shape4 = out_shape + + if expand: + out_shape1 = out_shape + out_shape2 = out_shape * 2 + out_shape3 = out_shape * 4 + if len(in_shape) >= 4: + out_shape4 = out_shape * 8 + + scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + if len(in_shape) >= 4: + scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups) + + return scratch + + +class ResidualConvUnit(nn.Module): + """Residual convolution module. + """ + + def __init__(self, features, activation, bn): + """Init. + + Args: + features (int): number of features + """ + super().__init__() + + self.bn = bn + + self.groups=1 + + self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups) + + self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups) + + if self.bn == True: + self.bn1 = nn.BatchNorm2d(features) + self.bn2 = nn.BatchNorm2d(features) + + self.activation = activation + + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x): + """Forward pass. + + Args: + x (tensor): input + + Returns: + tensor: output + """ + + out = self.activation(x) + out = self.conv1(out) + if self.bn == True: + out = self.bn1(out) + + out = self.activation(out) + out = self.conv2(out) + if self.bn == True: + out = self.bn2(out) + + if self.groups > 1: + out = self.conv_merge(out) + + return self.skip_add.add(out, x) + + +class FeatureFusionBlock(nn.Module): + """Feature fusion block. + """ + + def __init__( + self, + features, + activation, + deconv=False, + bn=False, + expand=False, + align_corners=True, + size=None + ): + """Init. + + Args: + features (int): number of features + """ + super(FeatureFusionBlock, self).__init__() + + self.deconv = deconv + self.align_corners = align_corners + + self.groups=1 + + self.expand = expand + out_features = features + if self.expand == True: + out_features = features // 2 + + self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1) + + self.resConfUnit1 = ResidualConvUnit(features, activation, bn) + self.resConfUnit2 = ResidualConvUnit(features, activation, bn) + + self.skip_add = nn.quantized.FloatFunctional() + + self.size=size + + def forward(self, *xs, size=None): + """Forward pass. + + Returns: + tensor: output + """ + output = xs[0] + + if len(xs) == 2: + res = self.resConfUnit1(xs[1]) + output = self.skip_add.add(output, res) + + output = self.resConfUnit2(output) + + if (size is None) and (self.size is None): + modifier = None #{"scale_factor": 2} + elif size is None: + modifier = {"size": self.size} + else: + modifier = {"size": size} + + if modifier: + output = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners) + + output = self.out_conv(output) + + return output diff --git a/cv/defom-stereo/depth_anything_v2/util/transform.py b/cv/defom-stereo/depth_anything_v2/util/transform.py new file mode 100644 index 0000000..ee1c0e7 --- /dev/null +++ b/cv/defom-stereo/depth_anything_v2/util/transform.py @@ -0,0 +1,158 @@ +import numpy as np +import cv2 + + +class Resize(object): + """Resize sample to given size (width, height). + """ + + def __init__( + self, + width, + height, + resize_target=True, + keep_aspect_ratio=False, + ensure_multiple_of=1, + resize_method="lower_bound", + image_interpolation_method=cv2.INTER_AREA, + ): + """Init. + + Args: + width (int): desired output width + height (int): desired output height + resize_target (bool, optional): + True: Resize the full sample (image, mask, target). + False: Resize image only. + Defaults to True. + keep_aspect_ratio (bool, optional): + True: Keep the aspect ratio of the input sample. + Output sample might not have the given width and height, and + resize behaviour depends on the parameter 'resize_method'. + Defaults to False. + ensure_multiple_of (int, optional): + Output width and height is constrained to be multiple of this parameter. + Defaults to 1. + resize_method (str, optional): + "lower_bound": Output will be at least as large as the given size. + "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.) + "minimal": Scale as least as possible. (Output size might be smaller than given size.) + Defaults to "lower_bound". + """ + self.__width = width + self.__height = height + + self.__resize_target = resize_target + self.__keep_aspect_ratio = keep_aspect_ratio + self.__multiple_of = ensure_multiple_of + self.__resize_method = resize_method + self.__image_interpolation_method = image_interpolation_method + + def constrain_to_multiple_of(self, x, min_val=0, max_val=None): + y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) + + if max_val is not None and y > max_val: + y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int) + + if y < min_val: + y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int) + + return y + + def get_size(self, width, height): + # determine new height and width + scale_height = self.__height / height + scale_width = self.__width / width + + if self.__keep_aspect_ratio: + if self.__resize_method == "lower_bound": + # scale such that output size is lower bound + if scale_width > scale_height: + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + elif self.__resize_method == "upper_bound": + # scale such that output size is upper bound + if scale_width < scale_height: + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + elif self.__resize_method == "minimal": + # scale as least as possbile + if abs(1 - scale_width) < abs(1 - scale_height): + # fit width + scale_height = scale_width + else: + # fit height + scale_width = scale_height + else: + raise ValueError(f"resize_method {self.__resize_method} not implemented") + + if self.__resize_method == "lower_bound": + new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height) + new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width) + elif self.__resize_method == "upper_bound": + new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height) + new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width) + elif self.__resize_method == "minimal": + new_height = self.constrain_to_multiple_of(scale_height * height) + new_width = self.constrain_to_multiple_of(scale_width * width) + else: + raise ValueError(f"resize_method {self.__resize_method} not implemented") + + return (new_width, new_height) + + def __call__(self, sample): + width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0]) + + # resize sample + sample["image"] = cv2.resize(sample["image"], (width, height), interpolation=self.__image_interpolation_method) + + if self.__resize_target: + if "depth" in sample: + sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST) + + if "mask" in sample: + sample["mask"] = cv2.resize(sample["mask"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST) + + return sample + + +class NormalizeImage(object): + """Normlize image by given mean and std. + """ + + def __init__(self, mean, std): + self.__mean = mean + self.__std = std + + def __call__(self, sample): + sample["image"] = (sample["image"] - self.__mean) / self.__std + + return sample + + +class PrepareForNet(object): + """Prepare sample for usage as network input. + """ + + def __init__(self): + pass + + def __call__(self, sample): + image = np.transpose(sample["image"], (2, 0, 1)) + sample["image"] = np.ascontiguousarray(image).astype(np.float32) + + if "depth" in sample: + depth = sample["depth"].astype(np.float32) + sample["depth"] = np.ascontiguousarray(depth) + + if "mask" in sample: + sample["mask"] = sample["mask"].astype(np.float32) + sample["mask"] = np.ascontiguousarray(sample["mask"]) + + return sample diff --git a/cv/defom-stereo/environment.yaml b/cv/defom-stereo/environment.yaml new file mode 100644 index 0000000..51fc8c8 --- /dev/null +++ b/cv/defom-stereo/environment.yaml @@ -0,0 +1,29 @@ +name: defomstereo +channels: + - pytorch + - nvidia + - xformers + - conda-forge + - bioconda + - defaults +dependencies: + - python=3.9 + - cudatoolkit=11.8.0 + - pytorch::pytorch=2.1.1 + - pytorch::pytorch-cuda=11.8.0 + - pytorch::torchvision=0.16.1 + - xformers::xformers=0.0.23 + - matplotlib + - tensorboard + - scipy + - opencv + - tqdm + - opt_einsum + - imageio + - scikit-image + - p7zip + - pip + - pip: + - gradio_imageslider + - gradio==4.29.0 + diff --git a/cv/defom-stereo/evaluate_stereo.py b/cv/defom-stereo/evaluate_stereo.py new file mode 100644 index 0000000..078515b --- /dev/null +++ b/cv/defom-stereo/evaluate_stereo.py @@ -0,0 +1,448 @@ +from __future__ import print_function, division +import sys + +import argparse +import time +import logging +import numpy as np +import torch +import torch.nn.functional as F +torch.cuda.empty_cache() +from PIL import Image + +from tqdm import tqdm +from core.defom_stereo import DEFOMStereo, autocast + +import core.stereo_datasets as datasets +from core.utils.utils import InputPadder + + +def count_parameters(model): + return sum(p.numel() for p in model.parameters()), sum(p.numel() for p in model.parameters() if p.requires_grad) + + +@torch.no_grad() +def validate_things(model, iters=32, scale_iters=8, mixed_prec=False, max_disp=192, bad_threshold=1.0): + """ Peform validation using the FlyingThings3D (TEST) split """ + model.eval() + val_dataset = datasets.SceneFlowDatasets(dstype='frames_finalpass', things_test=True) + + out_list, epe_list, elapsed_list = [], [], [] + for val_id in tqdm(range(len(val_dataset))): + data_blob = val_dataset[val_id] + image1 = data_blob["img1"][None].cuda() + image2 = data_blob["img2"][None].cuda() + disp_gt = data_blob["disp"] + valid = data_blob["valid"] + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + start = time.time() + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + end = time.time() + if val_id > 50: + elapsed_list.append(end-start) + + disp_pr = padder.unpad(disp_pr).cpu().squeeze(0) + assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape) + epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0) + + epe = epe.flatten() + val = (valid.flatten() >= 0.5) & (disp_gt.abs().flatten() < max_disp) + + if np.isnan(epe[val].mean().item()): + continue + out = (epe > bad_threshold) + image_out = out[val].float().mean().item() + image_epe = epe[val].mean().item() + if val_id < 9 or (val_id+1) % 10 == 0: + logging.info(f"Fhythings3D Iter {val_id+1} out of {len(val_dataset)}. EPE {round(image_epe,4)} Out{bad_threshold} {round(image_out,4)}. Runtime: {format(end-start, '.3f')}s ({format(1/(end-start), '.2f')}-FPS)") + + epe_list.append(image_epe) + out_list.append(out[val].cpu().numpy()) + + epe_list = np.array(epe_list) + out_list = np.concatenate(out_list) + + epe = np.mean(epe_list) + out = 100 * np.mean(out_list) + avg_runtime = np.mean(elapsed_list) + + print(f"Validation FlyingThings: EPE {epe}, Out{bad_threshold} {out}, " + f"{format(1/avg_runtime, '.2f')}-FPS ({format(avg_runtime, '.3f')}s)") + return {'things-epe': epe, 'things-out': out} + + +@torch.no_grad() +def validate_eth3d(model, iters=32, scale_iters=8, mixed_prec=False): + """ Peform validation using the ETH3D (train) split """ + model.eval() + aug_params = {} + val_dataset = datasets.ETH3D(aug_params, is_eval=True) + + out_list, epe_list = [], [] + for val_id in tqdm(range(len(val_dataset))): + data_blob = val_dataset[val_id] + image1 = data_blob["img1"][None].cuda() + image2 = data_blob["img2"][None].cuda() + disp_gt = data_blob["disp"] + valid = data_blob["valid"] + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + disp_pr = padder.unpad(disp_pr).cpu().squeeze(0) + assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape) + epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0) + + epe_flattened = epe.flatten() + val = valid.flatten() >= 0.5 + out = (epe_flattened > 1.0) + image_out = out[val].float().mean().item() + image_epe = epe_flattened[val].mean().item() + logging.info(f"ETH3D {val_id+1} out of {len(val_dataset)}. EPE {round(image_epe,4)} D1 {round(image_out,4)}") + epe_list.append(image_epe) + out_list.append(image_out) + + epe_list = np.array(epe_list) + out_list = np.array(out_list) + + epe = np.mean(epe_list) + out1 = 100 * np.mean(out_list) + + print("Validation ETH3D: EPE %f, Out1 %f" % (epe, out1)) + return {'eth3d-epe': epe, 'eth3d-out1': out1} + + +@torch.no_grad() +def validate_kitti(model, iters=32, scale_iters=8, split='15', mixed_prec=False): + """ Peform validation using the KITTI-2015/2012 (train) split """ + model.eval() + aug_params = {} + val_dataset = datasets.KITTI(aug_params, split=split, image_set='training', is_eval=True) + torch.backends.cudnn.benchmark = True + + out_list, epe_list, elapsed_list = [], [], [] + for val_id in range(len(val_dataset)): + data_blob = val_dataset[val_id] + image1 = data_blob["img1"][None].cuda() + image2 = data_blob["img2"][None].cuda() + disp_gt = data_blob["disp"] + valid = data_blob["valid"] + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + start = time.time() + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + end = time.time() + if val_id > 50: + elapsed_list.append(end-start) + + disp_pr = padder.unpad(disp_pr).cpu().squeeze(0) + assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape) + epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0) + + epe_flattened = epe.flatten() + val = valid.flatten() >= 0.5 + + out = (epe_flattened > 3.0) + image_out = out[val].float().mean().item() + image_epe = epe_flattened[val].mean().item() + if val_id < 9 or (val_id+1) % 10 == 0: + logging.info(f"KITTI{split} Iter {val_id+1} out of {len(val_dataset)}. EPE {round(image_epe,4)} Out3 {round(image_out,4)}. Runtime: {format(end-start, '.3f')}s ({format(1/(end-start), '.2f')}-FPS)") + epe_list.append(epe_flattened[val].mean().item()) + out_list.append(out[val].cpu().numpy()) + + epe_list = np.array(epe_list) + out_list = np.concatenate(out_list) + + epe = np.mean(epe_list) + out3 = 100 * np.mean(out_list) + + avg_runtime = np.mean(elapsed_list) + + print(f"Validation KITTI{split}: EPE {epe}, Out3 {out3}, " + f"{format(1/avg_runtime, '.2f')}-FPS ({format(avg_runtime, '.3f')}s)") + return {f'kitti{split}-epe': epe, f'kitti{split}-out3': out3} + + +@torch.no_grad() +def validate_middlebury(model, iters=32, scale_iters=8, split='H', mixed_prec=False): + """ Peform validation using the Middlebury-V3 dataset """ + model.eval() + aug_params = {} + val_dataset = datasets.Middlebury(aug_params, split=split, is_eval=True) + + out_list, epe_list = [], [] + for val_id in range(len(val_dataset)): + data_blob = val_dataset[val_id] + image1 = data_blob["img1"][None].cuda() + image2 = data_blob["img2"][None].cuda() + disp_gt = data_blob["disp"] + valid = data_blob["valid"] + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + disp_pr = padder.unpad(disp_pr).cpu().squeeze(0) + assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape) + epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0) + + epe_flattened = epe.flatten() + val = (valid.reshape(-1) >= 0.5) & (disp_gt.reshape(-1) < 1000) + + out = (epe_flattened > 2.0) + image_out = out[val].float().mean().item() + image_epe = epe_flattened[val].mean().item() + logging.info(f"Middlebury Iter {val_id+1} out of {len(val_dataset)}. " + f"EPE {round(image_epe,4)} Out2 {round(image_out,4)}") + epe_list.append(image_epe) + out_list.append(image_out) + + epe_list = np.array(epe_list) + out_list = np.array(out_list) + + epe = np.mean(epe_list) + out2 = 100 * np.mean(out_list) + + print(f"Validation Middlebury{split}: EPE {epe}, Out2 {out2}") + return {f'middlebury{split}-epe': epe, f'middlebury{split}-out2': out2} + + +def compute_nontexture(x, weight=None, c1=0.01**2, c2=0.03**2, weight_epsilon=0.01, window=33, threshold=0.95, split="F"): + + if split=="H": + scale = 2 + threshold += 0.02 + elif split=="Q": + scale = 4 + threshold += 0.03 + else: + scale = 1 + + x = F.interpolate(x, scale_factor=scale, mode='bilinear', align_corners=True) + + if x.max()>1: + x = x/x.max() + + y = F.pad(x, (1, 1, 1, 1), mode='replicate') + _, _, h, w = y.shape + #y = y[..., 0:h-2, 1:w-1] #(y[..., 0:h-2, 1:w-1] + y[..., 2:h, 1:w-1] + y[..., 1:h-1, 0:w-2] + y[..., 1:h-1, 2:w])/4.0 + + x = F.pad(x, (window//2, window//2, window//2, window//2), mode='replicate') + if c1 == float('inf') and c2 == float('inf'): + raise ValueError( + 'Both c1 and c2 are infinite, SSIM loss is zero. This is ' + 'likely unintended.') + _, _, H, W = x.shape + + if weight is None: + weight = torch.ones((H, W)).to(x) + else: + assert weight.shape == (H, W), \ + f'image shape is {(H, W)}, but weight shape is {weight.shape}' + weight = weight[None, None, ...] + average_pooled_weight = F.avg_pool2d(weight, (window, window), stride=(1, 1)) + weight_plus_epsilon = weight + weight_epsilon + inverse_average_pooled_weight = 1.0 / ( + average_pooled_weight + weight_epsilon) + + def weighted_avg_pool(z): + weighted_avg = F.avg_pool2d( + z * weight_plus_epsilon, (window, window), stride=(1, 1)) + return weighted_avg * inverse_average_pooled_weight + + mu_x = weighted_avg_pool(x) + sigma_x = weighted_avg_pool(x**2) - mu_x**2 + + def ssim(x, y): + y = F.pad(y, (window//2, window//2, window//2, window//2), mode='replicate') + mu_y = weighted_avg_pool(y) + sigma_y = weighted_avg_pool(y**2) - mu_y**2 + sigma_xy = weighted_avg_pool(x * y) - mu_x * mu_y + if c1 == float('inf'): + ssim_n = (2 * sigma_xy + c2) + ssim_d = (sigma_x + sigma_y + c2) + elif c2 == float('inf'): + ssim_n = 2 * mu_x * mu_y + c1 + ssim_d = mu_x**2 + mu_y**2 + c1 + else: + ssim_n = (2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2) + ssim_d = (mu_x**2 + mu_y**2 + c1) * (sigma_x + sigma_y + c2) + + result = ssim_n / ssim_d + + result = F.avg_pool2d(result, (scale, scale), stride=(scale, scale)) + + return result + + mask = (ssim(x, y[..., 0:h-2, 1:w-1])>threshold) & (ssim(x, y[..., 2:h, 1:w-1])>threshold) & (ssim(x, y[..., 1:h-1, 0:w-2])>threshold) & (ssim(x, y[..., 1:h-1, 2:w])>threshold) + mask = mask[0, 0] & mask[0, 1] & mask[0, 2] + + return mask.cpu().numpy() + + +@torch.no_grad() +def validate_middlebury_indetail(model, iters=32, scale_iters=8, split='H', mixed_prec=False): + """ Peform validation using the Middlebury-V3 dataset """ + model.eval() + aug_params = {} + val_dataset = datasets.Middlebury(aug_params, split=split, is_eval=True) + + out_list, epe_list, portion_list = [[], [], [], []], [[], [], [], []], [[], [], [], []] + for val_id in range(len(val_dataset)): + data_blob = val_dataset[val_id] + image1 = data_blob["img1"][None].cuda() + image2 = data_blob["img2"][None].cuda() + disp_gt = data_blob["disp"] + valid = data_blob["valid"] + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + disp_pr = padder.unpad(disp_pr).cpu().squeeze(0) + assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape) + epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0) + + epe_flattened = epe.flatten() + + occ_mask = Image.open(data_blob["imageL_file"].replace('im0.png', 'mask0nocc.png')).convert('L') + occ_mask = np.ascontiguousarray(occ_mask, dtype=np.float32).flatten() + val_all = (valid.reshape(-1) >= 0.5) & (disp_gt.reshape(-1) < 1000) + val_occ = val_all & (occ_mask==128) + val_nocc = val_all & (occ_mask==255) + + val_ntt = val_all & compute_nontexture(data_blob["img1"][None].cuda(), split=split).flatten() + + out = (epe_flattened > 2.0) + image_out = out[val_all].float().mean().item() + image_epe = epe_flattened[val_all].mean().item() + + image_out_occ = out[val_occ].float().mean().item() + image_epe_occ = epe_flattened[val_occ].mean().item() + + image_out_nocc = out[val_nocc].float().mean().item() + image_epe_nocc = epe_flattened[val_nocc].mean().item() + + image_out_ntt = out[val_ntt].float().mean().item() + image_epe_ntt = epe_flattened[val_ntt].mean().item() + + logging.info(f"Middlebury Iter {val_id+1} out of {len(val_dataset)}. " + f"All({round((val_all.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe,4)} Out2 {round(image_out,4)}, \n " + f"Occ({round((val_occ.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe_occ,4)} Out2 {round(image_out_occ,4)}, " + f"NOcc({round((val_nocc.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe_nocc,4)} Out2 {round(image_out_nocc,4)}, " + f"NonTexture({round((val_ntt.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe_ntt,4)} Out2 {round(image_out_ntt,4)}") + + epe_list[0].append(image_epe) + out_list[0].append(image_out) + portion_list[0].append((val_all.sum()/val_all.sum()).item()) + epe_list[1].append(image_epe_occ) + out_list[1].append(image_out_occ) + portion_list[1].append((val_occ.sum()/val_all.sum()).item()) + epe_list[2].append(image_epe_nocc) + out_list[2].append(image_out_nocc) + portion_list[2].append((val_nocc.sum()/val_all.sum()).item()) + epe_list[3].append(image_epe_ntt) + out_list[3].append(image_out_ntt) + portion_list[3].append((val_ntt.sum()/val_all.sum()).item()) + + epe_list = np.array(epe_list) + out_list = np.array(out_list) + portion_list = np.array(portion_list) + + epe = np.mean(epe_list, axis=1) + out2 = 100 * np.mean(out_list, axis=1) + portion = 100 * np.mean(portion_list, axis=1) + + print(f"Validation Middlebury{split}: All({round(portion[0],8)}%): EPE {round(epe[0],8)} Out2 {round(out2[0],8)}, \n" + f"Occ({round(portion[1],8)}%): EPE {round(epe[1],8)} Out2 {round(out2[1],8)}, " + f"NOcc({round(portion[2],8)}%): EPE {round(epe[2],8)} Out2 {round(out2[2],8)}, " + f"NonTexture({round(portion[3],8)}%): EPE {round(epe[3],8)} Out2 {round(out2[3],8)}") + return {f'middlebury{split}-epe': epe[0], f'middlebury{split}-out2': out2[0]} + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--restore_ckpt', help="restore checkpoint", default=None) + parser.add_argument('--datasets', nargs='+', type=str, help="dataset for evaluation", default=["things"], + choices=["things", "eth3d", "kitti12", "kitti15"] + [f"middlebury_{s}" for s in 'FHQ']) + parser.add_argument('--indetail', action='store_true', help='evaluate middlebury in detail (for different regions)') + + parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') + parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during forward pass') + parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.") + + # Architecure choices + parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg']) + parser.add_argument('--idepth_scale', type=float, default=0.5, help="the scale of inverse depth to initialize disparity") + parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions") + parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation") + parser.add_argument('--shared_backbone', action='store_true', help="use a single backbone for the context and feature encoders") + parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid") + parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid") + parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0], + help='the list of scaling factors of disparity') + parser.add_argument('--scale_corr_radius', type=int, default=2, + help="width of the correlation pyramid for scaled disparity") + + parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], help="resolution of the disparity field (1/2^K)") + parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder") + parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels") + + args = parser.parse_args() + + model = DEFOMStereo(args) + + logging.basicConfig(level=logging.INFO, + format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s') + + if args.restore_ckpt is not None: + assert args.restore_ckpt.endswith(".pth") + logging.info("Loading checkpoint...") + checkpoint = torch.load(args.restore_ckpt, map_location='cuda') + if 'model' in checkpoint: + model.load_state_dict(checkpoint['model']) + else: + model.load_state_dict(checkpoint) + logging.info(f"Done loading checkpoint") + + model.cuda() + model.eval() + + print(f"The model has {format(count_parameters(model)[1]/1e6, '.2f')}M learnable parameters.") + + # The CUDA implementations of the correlation volume prevent half-precision + # rounding errors in the correlation lookup. This allows us to use mixed precision + # in the entire forward pass, not just in the GRUs & feature extractors. + use_mixed_precision = args.corr_implementation.endswith("_cuda") + + if 'things' in args.datasets: + validate_things(model, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision) + + if 'eth3d' in args.datasets: + validate_eth3d(model, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision) + + if 'kitti12' in args.datasets: + validate_kitti(model, iters=args.valid_iters, scale_iters=args.scale_iters, split='12', mixed_prec=use_mixed_precision) + + if 'kitti15' in args.datasets: + validate_kitti(model, iters=args.valid_iters, scale_iters=args.scale_iters, split='15', mixed_prec=use_mixed_precision) + + for s in 'FHQ': + if f"middlebury_{s}" in args.datasets: + if args.indetail: + validate_middlebury_indetail(model, iters=args.valid_iters, scale_iters=args.scale_iters, split=s, mixed_prec=use_mixed_precision) + else: + validate_middlebury(model, iters=args.valid_iters, scale_iters=args.scale_iters, split=s, mixed_prec=use_mixed_precision) + diff --git a/cv/defom-stereo/make_submission.py b/cv/defom-stereo/make_submission.py new file mode 100644 index 0000000..7772ce9 --- /dev/null +++ b/cv/defom-stereo/make_submission.py @@ -0,0 +1,237 @@ +from __future__ import print_function, division + +import argparse +import logging +import numpy as np +import torch +from tqdm import tqdm +import time +import os +import cv2 +import sys + + +from core.defom_stereo import DEFOMStereo, autocast + +import core.stereo_datasets as datasets +from core.utils.utils import InputPadder +from core.utils.frame_utils import writePFM + + +def makedirs(path): + if not os.path.exists(path): + os.makedirs(path) + + +def StrToBytes(text): + if sys.version_info[0] == 2: + return text + else: + return bytes(text, 'UTF-8') + + +def count_parameters(model): + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +@torch.no_grad() +def test_eth3d(model, save_path, iters=32, scale_iters=3, mixed_prec=False): + """ Peform validation using the ETH3D (train) split """ + model.eval() + aug_params = {} + + test_dataset = datasets.ETH3D(aug_params, split='testing', is_test=True) + training_dataset = datasets.ETH3D(aug_params, split='training', is_test=True) + dataset = test_dataset + training_dataset + torch.backends.cudnn.benchmark = True + + for test_id in tqdm(range(len(dataset))): + img1, img2, imageL_file = dataset[test_id] + image1 = img1[None].cuda() + image2 = img2[None].cuda() + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + start = time.time() + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + end = time.time() + runtime = end - start + disp = padder.unpad(disp_pr).cpu().squeeze().numpy() + disp[disp < 0] = 0 + disp[disp > 64] = 64 + + names = imageL_file.split("/") + save_sub_path = os.path.join(save_path, "low_res_"+names[-3]) + makedirs(save_sub_path) + + disp_path = os.path.join(save_sub_path, names[-2] + '.pfm') + writePFM(disp_path, disp) + + txt_path = os.path.join(save_sub_path, names[-2] + '.txt') + with open(txt_path, 'wb') as time_file: + time_file.write(StrToBytes('runtime ' + str(runtime))) + + +@torch.no_grad() +def test_kitti(model, save_path, iters=32, scale_iters=3, split='15', mixed_prec=False): + """ Peform testing on the KITTI-2015 (test) split """ + model.eval() + aug_params = {} + save_path = os.path.join(save_path, "disp_0") + makedirs(save_path) + + test_dataset = datasets.KITTI(aug_params, split=split, image_set='testing', is_test=True) + + runtime_sum = 0.0 + runtime_count = 0 + + for test_id in tqdm(range(len(test_dataset))): + img1, img2, imageL_file = test_dataset[test_id] + image1 = img1[None].cuda() + image2 = img2[None].cuda() + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + start = time.time() + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + end = time.time() + runtime = end - start + runtime_sum += runtime + runtime_count += 1 + + disp = padder.unpad(disp_pr).cpu().squeeze().numpy() + disp[disp < 0] = 0 + disp[disp > 240] = 240 + disp = np.uint16(disp*256) + + name = imageL_file.split('/')[-1] + path = os.path.join(save_path, name) + cv2.imwrite(path, disp, [cv2.IMWRITE_PNG_COMPRESSION, 9]) + + print('The average runtime on Kitti test images is (you will need this for the submission): ' + + str(runtime_sum / runtime_count) + " seconds") + + +@torch.no_grad() +def test_middlebury(model, save_path, iters=32, scale_iters=8, split='F', mixed_prec=False, method_name="DEFOM-Stereo"): + """ Peform validation using the Middlebury-V3 dataset """ + model.eval() + aug_params = {} + test_dataset = datasets.Middlebury(aug_params, split=split, image_set='test', is_test=True) + training_dataset = datasets.Middlebury(aug_params, split=split, image_set='training', is_test=True) + dataset = test_dataset + training_dataset + torch.backends.cudnn.benchmark = True + + for test_id in tqdm(range(len(dataset))): + img1, img2, imageL_file = dataset[test_id] + image1 = img1[None].cuda() + image2 = img2[None].cuda() + + padder = InputPadder(image1.shape, divis_by=32) + image1, image2 = padder.pad(image1, image2) + + with autocast(enabled=mixed_prec): + start = time.time() + disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True) + end = time.time() + runtime = end - start + disp = padder.unpad(disp_pr).cpu().squeeze().numpy() + disp[disp < 0] = 0 + disp[disp > 800] = 800 + + names = imageL_file.split("/") + save_sub_path = os.path.join(save_path, names[-3], names[-2]) + makedirs(save_sub_path) + + disp_path = os.path.join(save_sub_path, 'disp0' + method_name + '.pfm') + writePFM(disp_path, disp) + + txt_path = os.path.join(save_sub_path, 'time' + method_name + '.txt') + with open(txt_path, 'wb') as time_file: + time_file.write(StrToBytes(str(runtime))) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + + parser.add_argument('--restore_ckpt', help="restore checkpoint", default=None) + parser.add_argument('--datasets', nargs='+', type=str, help="dataset for evaluation", default=["kitti12", "kitti15"], + choices=["eth3d", "kitti12", "kitti15"] + [f"middlebury_{s}" for s in 'FHQ']) + parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') + parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during forward pass') + parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.") + parser.add_argument('--method_name', default="DEFOM-Stereo", help="the method to test") + + # Architecure choices + parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg']) + parser.add_argument('--idepth_scale', type=float, default=0.5, + help="the scale of inverse depth to initialize disparity") + parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128] * 3, + help="hidden state and context dimensions") + parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", + help="correlation volume implementation") + parser.add_argument('--shared_backbone', action='store_true', + help="use a single backbone for the context and feature encoders") + parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid") + parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid") + parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0], + help='the list of scaling factors of disparity') + parser.add_argument('--scale_corr_radius', type=int, default=2, + help="width of the correlation pyramid for scaled disparity") + + parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], + help="resolution of the disparity field (1/2^K)") + parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], + help="normalization of context encoder") + parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels") + + args = parser.parse_args() + + model = DEFOMStereo(args) + + logging.basicConfig(level=logging.INFO, + format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s') + + if args.restore_ckpt is not None: + assert args.restore_ckpt.endswith(".pth") + logging.info("Loading checkpoint...") + checkpoint = torch.load(args.restore_ckpt, map_location='cuda') + model.load_state_dict(checkpoint, strict=True) + logging.info(f"Done loading checkpoint") + + model.cuda() + model.eval() + + print(f"The model has {format(count_parameters(model)/1e6, '.2f')}M learnable parameters.") + + # The CUDA implementations of the correlation volume prevent half-precision + # rounding errors in the correlation lookup. This allows us to use mixed precision + # in the entire forward pass, not just in the GRUs & feature extractors. + use_mixed_precision = args.corr_implementation.endswith("_cuda") + + if 'eth3d' in args.datasets: + save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "eth3d" + makedirs(save_path) + test_eth3d(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision) + + if 'kitti12' in args.datasets: + save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "kitti12" + makedirs(save_path) + test_kitti(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision, split='12') + + if 'kitti15' in args.datasets: + save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "kitti15" + makedirs(save_path) + test_kitti(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision, split='15') + + for s in 'FHQ': + if f"middlebury_{s}" in args.datasets: + save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + f"middlebury_{s}" + makedirs(save_path) + test_middlebury(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, split=s, + method_name=args.method_name, mixed_prec=use_mixed_precision) + diff --git a/cv/defom-stereo/requirements.txt b/cv/defom-stereo/requirements.txt new file mode 100644 index 0000000..58faaec --- /dev/null +++ b/cv/defom-stereo/requirements.txt @@ -0,0 +1,16 @@ +--extra-index-url https://download.pytorch.org/whl/cu118 +torch==2.1.1 +torchvision==0.16.1 +xformers==0.0.23 +gradio_imageslider +gradio==4.29.0 +matplotlib +tensorboard +scipy +tqdm +opt_einsum +imageio +scikit-image +pillow +timm +gdown diff --git a/cv/defom-stereo/scripts/download_dav2.sh b/cv/defom-stereo/scripts/download_dav2.sh new file mode 100644 index 0000000..54a9716 --- /dev/null +++ b/cv/defom-stereo/scripts/download_dav2.sh @@ -0,0 +1,6 @@ +cd checkpoints +wget https://huggingface.co/depth-anything/Depth-Anything-V2-Small/resolve/main/depth_anything_v2_vits.pth +wget https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth +cd .. + + diff --git a/cv/defom-stereo/scripts/download_models.sh b/cv/defom-stereo/scripts/download_models.sh new file mode 100644 index 0000000..adb9dd8 --- /dev/null +++ b/cv/defom-stereo/scripts/download_models.sh @@ -0,0 +1,10 @@ +cd checkpoints +gdown https://drive.google.com/uc?id=1XuAM4vqzura_6NKN70hMW5lFD4TafnDL +gdown https://drive.google.com/uc?id=1FNt-SDysG5bUOmjZ91mzH2V_TXvLCvr5 +gdown https://drive.google.com/uc?id=1qyXKO-Nxq3ndl2H0deQpo6BSvwlGKYEg +gdown https://drive.google.com/uc?id=1Dy1eGDdtkp2GQYQRTvMwR-3eAzaRCe_k +gdown https://drive.google.com/uc?id=1duHLtUCDNIA76m6Fqwa7hv-aBMY-P3mg +gdown https://drive.google.com/uc?id=1xEPS7gceJSFn_IHdzebBCgQaRNGwf1aG +cd .. + + diff --git a/cv/defom-stereo/scripts/evaluate.sh b/cv/defom-stereo/scripts/evaluate.sh new file mode 100644 index 0000000..c6cbb0d --- /dev/null +++ b/cv/defom-stereo/scripts/evaluate.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash + +# evalutate on scene flow +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets things \ +--dinov2_encoder vitl + +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets things \ +--dinov2_encoder vits + +# evalutate on kitti12, kitti15, and eth3d +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets kitti12 kitti15 eth3d \ +--dinov2_encoder vitl + +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets kitti12 kitti15 eth3d \ +--dinov2_encoder vits + + +# evalutate on Middlebury; when evaluating defomstereo_vitl on Middlebury_F +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets middlebury_F middlebury_H middlebury_Q \ +--dinov2_encoder vitl + +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets middlebury_F middlebury_H middlebury_Q \ +--dinov2_encoder vits + +# evalutate on different region. +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets middlebury_F middlebury_H middlebury_Q \ +--indetail \ +--dinov2_encoder vitl + +python evaluate_stereo.py \ +--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets middlebury_F middlebury_H middlebury_Q \ +--indetail \ +--dinov2_encoder vits + diff --git a/cv/defom-stereo/scripts/make_submission.sh b/cv/defom-stereo/scripts/make_submission.sh new file mode 100644 index 0000000..edeadc0 --- /dev/null +++ b/cv/defom-stereo/scripts/make_submission.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash + +# make submission to stereo benchmarks + +# make submission for kitti12 and kitti15 +python make_submission.py \ +--restore_ckpt checkpoints/defomstereo_vitl_kitti.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets kitti12 kitti15 \ +--dinov2_encoder vitl + +# make submission for eth3d +python make_submission.py \ +--restore_ckpt checkpoints/defomstereo_vitl_eth3d.pth \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets eth3d \ +--dinov2_encoder vitl + +# make submission for middlebury +python make_submission.py \ +--restore_ckpt checkpoints/defomstereo_vitl_middlebury.pth \ +--method_name DEFOM-Stereo \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets middlebury_F \ +--dinov2_encoder vitl + +# make submission for kitti15, middlebury and eth3d using the RVC model +python make_submission.py \ +--restore_ckpt checkpoints/defomstereo_vits_rvc.pth \ +--method_name DEFOM-Stereo_RVC \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--datasets kitti15 middlebury_F eth3d \ +--dinov2_encoder vits diff --git a/cv/defom-stereo/scripts/train_eth3d.sh b/cv/defom-stereo/scripts/train_eth3d.sh new file mode 100644 index 0000000..692830e --- /dev/null +++ b/cv/defom-stereo/scripts/train_eth3d.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash + +# trained on 4 x 24GB 3090/4090 GPUs + +CHECKPOINT_DIR=checkpoints/defomstereo_vitl_eth3d_pretrain && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9994 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vitl_eth3d_pretrain \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets tartan_air sceneflow sintel_stereo eth3d instereo2k crestereo \ +--train_folds 1 1 50 1000 100 2 \ +--num_steps 300000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vitl \ +--image_size 384 512 \ +--resume_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \ + +CHECKPOINT_DIR=checkpoints/defomstereo_vitl_eth3d && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9993 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vitl_eth3d \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets eth3d instereo2k crestereo \ +--train_folds 1000 10 1 \ +--num_steps 90000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vitl \ +--image_size 384 512 \ +--resume_ckpt checkpoints/defomstereo_vitl_eth3d_pretrain.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log diff --git a/cv/defom-stereo/scripts/train_kitti.sh b/cv/defom-stereo/scripts/train_kitti.sh new file mode 100644 index 0000000..961a140 --- /dev/null +++ b/cv/defom-stereo/scripts/train_kitti.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash + +# trained on 8 x 24GB 3090/4090 GPUs + +CHECKPOINT_DIR=checkpoints/defomstereo_vitl_kitti && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=8 --master_port=9992 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 4 5 6 7 \ +--name defomstereo_vitl_kitti \ +--batch_size 8 \ +--num_workers 4 \ +--train_datasets kitti12 kitti15 vkitti2 \ +--train_folds 50 50 1 \ +--num_steps 50000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vitl \ +--image_size 352 1216 \ +--resume_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log + + diff --git a/cv/defom-stereo/scripts/train_middlebury.sh b/cv/defom-stereo/scripts/train_middlebury.sh new file mode 100644 index 0000000..17af88d --- /dev/null +++ b/cv/defom-stereo/scripts/train_middlebury.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash + +# trained on 4 x 24GB 3090/4090 GPUs +CHECKPOINT_DIR=checkpoints/defomstereo_vitl_middlebury_pretrain && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9993 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vitl_middlebury_pretrain \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets tartan_air sceneflow falling_things instereo2k carla_highres crestereo middlebury_2014 middlebury_2021 middlebury_H \ +--train_folds 1 1 1 50 50 1 200 200 200 \ +--num_steps 200000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vitl \ +--image_size 384 512 \ +--resume_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \ + +# trained on 8 x 24GB 3090/4090 GPUs +CHECKPOINT_DIR=checkpoints/defomstereo_vitl_middlebury && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=8 --master_port=9993 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 4 5 6 7 \ +--name defomstereo_vitl_middlebury \ +--batch_size 8 \ +--num_workers 4 \ +--train_datasets crestereo instereo2k carla_highres middlebury_2014 middlebury_2021 middlebury_H middlebury_F falling_things \ +--train_folds 1 50 50 200 200 200 200 5 \ +--num_steps 100000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vitl \ +--image_size 512 768 \ +--resume_ckpt checkpoints/defomstereo_vitl_middlebury_pretrain.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log diff --git a/cv/defom-stereo/scripts/train_rvc.sh b/cv/defom-stereo/scripts/train_rvc.sh new file mode 100644 index 0000000..68e5fd5 --- /dev/null +++ b/cv/defom-stereo/scripts/train_rvc.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash + +# trained on 4 x 24GB 3090/4090 GPUs +CHECKPOINT_DIR=checkpoints/defomstereo_vits_rvc_pretrain && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9995 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vits_rvc_pretrain \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets tartan_air sceneflow irs 3dkenburns crestereo falling_things sintel_stereo vkitti2 carla_highres \ +--train_folds 1 1 1 1 1 1 3 3 80 \ +--num_steps 200000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vits \ +--image_size 384 768 \ +--resume_ckpt checkpoints/defomstereo_vits_sceneflow.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \ + +# trained on 4 x 24GB 3090/4090 GPUs +CHECKPOINT_DIR=checkpoints/defomstereo_vits_rvc_pretrain2 && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9996 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vits_rvc_pretrain2 \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets tartan_air irs 3dkenburns crestereo vkitti2 carla_highres kitti12 kitti15 middlebury_2005 middlebury_2006 middlebury_2014 middlebury_2021 middlebury_Q middlebury_H eth3d instereo2k booster \ +--train_folds 1 1 1 1 3 30 100 100 200 200 200 200 200 200 1000 20 10 \ +--num_steps 100000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vits \ +--image_size 384 768 \ +--resume_ckpt checkpoints/defomstereo_vits_rvc_pretrain.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \ + +# trained on 4 x 24GB 3090/4090 GPUs +CHECKPOINT_DIR=checkpoints/defomstereo_vits_rvc && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9997 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vits_rvc \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets tartan_air irs 3dkenburns crestereo vkitti2 carla_highres kitti12 kitti15 middlebury_2005 middlebury_2006 middlebury_2014 middlebury_2021 middlebury_Q middlebury_H eth3d instereo2k booster \ +--train_folds 1 1 1 1 3 30 2500 2500 200 200 200 200 200 200 1000 20 10 \ +--num_steps 20000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vits \ +--image_size 384 768 \ +--resume_ckpt checkpoints/defomstereo_vits_rvc_pretrain2.pth \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log diff --git a/cv/defom-stereo/scripts/train_sceneflow_vitl.sh b/cv/defom-stereo/scripts/train_sceneflow_vitl.sh new file mode 100644 index 0000000..8baab02 --- /dev/null +++ b/cv/defom-stereo/scripts/train_sceneflow_vitl.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash + +# trained on 4 x 24GB 3090/4090 GPUs + +CHECKPOINT_DIR=checkpoints/defomstereo_vitl_sceneflow && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=4 --master_port=9991 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 2 3 \ +--name defomstereo_vitl_sceneflow \ +--batch_size 8 \ +--num_workers 8 \ +--train_datasets sceneflow \ +--train_folds 1 \ +--num_steps 200000 \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vitl \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log + + diff --git a/cv/defom-stereo/scripts/train_sceneflow_vits.sh b/cv/defom-stereo/scripts/train_sceneflow_vits.sh new file mode 100644 index 0000000..ac2fc45 --- /dev/null +++ b/cv/defom-stereo/scripts/train_sceneflow_vits.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash + +# trained on 2 x 24GB 3090/4090 GPUs + +CHECKPOINT_DIR=checkpoints/defomstereo_vits_sceneflow && \ +mkdir -p ${CHECKPOINT_DIR} && \ +python -m torch.distributed.launch --nproc_per_node=2 --master_port=9990 train_stereo.py \ +--distributed \ +--launcher pytorch \ +--gpu_ids 0 1 \ +--name defomstereo_vits_sceneflow \ +--batch_size 8 \ +--num_workers 16 \ +--train_datasets sceneflow \ +--train_folds 1 \ +--num_steps 200000 \ +--mixed_precision \ +--n_downsample 2 \ +--train_iters 18 \ +--scale_iters 8 \ +--idepth_scale 0.5 \ +--corr_levels 2 \ +--corr_radius 4 \ +--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \ +--scale_corr_radius 2 \ +--dinov2_encoder vits \ +2>&1 | tee -a ${CHECKPOINT_DIR}/train.log + + diff --git a/cv/defom-stereo/train_stereo.py b/cv/defom-stereo/train_stereo.py new file mode 100644 index 0000000..8a08082 --- /dev/null +++ b/cv/defom-stereo/train_stereo.py @@ -0,0 +1,292 @@ +from __future__ import print_function, division +import os +import sys +import logging +import argparse +import numpy as np +from pathlib import Path +from tqdm import tqdm + +from torch.utils.tensorboard import SummaryWriter +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader + +from utils.dist_utils import get_dist_info, init_dist, setup_for_distributed +from utils.utils import * +from core.defom_stereo import DEFOMStereo + +from evaluate_stereo import validate_things, count_parameters +import core.stereo_datasets as datasets + +try: + from torch.cuda.amp import GradScaler +except: + # dummy GradScaler for PyTorch < 1.6 + class GradScaler: + def __init__(self): + pass + def scale(self, loss): + return loss + def unscale_(self, optimizer): + pass + def step(self, optimizer): + optimizer.step() + def update(self): + pass + + +def train(args): + + seed_everything(args.seed) + + if args.launcher == 'none': + args.distributed = False + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + else: + args.distributed = True + + # adjust batch size for each gpu + assert args.batch_size % torch.cuda.device_count() == 0 + args.batch_size = args.batch_size // torch.cuda.device_count() + + dist_params = dict(backend='nccl') + init_dist(args.launcher, **dist_params) + # re-set gpu_ids with distributed training mode + _, world_size = get_dist_info() + args.gpu_ids = range(world_size) + device = torch.device('cuda:{}'.format(args.local_rank)) + + setup_for_distributed(args.local_rank == 0) + + model = DEFOMStereo(args).to(device) + print("Parameter Count: %d, Trainable: %d" % count_parameters(model)) + + if args.distributed: + process_group = torch.distributed.new_group(list(range(len(args.gpu_ids)))) + model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group) + model = torch.nn.parallel.DistributedDataParallel( + model.to(device), + device_ids=[args.local_rank], + output_device=args.local_rank, + find_unused_parameters=True) + model_without_ddp = model.module + else: + if torch.cuda.device_count() > 1: + print('Use %d GPUs' % torch.cuda.device_count()) + model = torch.nn.DataParallel(model) + model_without_ddp = model.module + else: + model_without_ddp = model + + model_without_ddp.freeze_bn() # BatchNorm kept frozen if not distributed + + start_epoch = 0 + start_step = 0 + optimizer, scheduler = fetch_optimizer(args, model) + + if args.resume_ckpt: + assert args.resume_ckpt.endswith(".pth") + logging.info("Loading checkpoint: %s" % args.resume_ckpt) + loc = 'cuda:{}'.format(args.local_rank) if torch.cuda.is_available() else 'cpu' + checkpoint = torch.load(args.resume_ckpt, map_location=loc) + if 'model' in checkpoint: + model_without_ddp.load_state_dict(checkpoint['model'], strict=args.strict_resume) + else: + model_without_ddp.load_state_dict(checkpoint, strict=args.strict_resume) + + if 'optimizer' in checkpoint and 'step' in checkpoint and 'epoch' in checkpoint and not \ + args.no_resume_optimizer: + print('Load optimizer') + start_step = checkpoint['step'] + start_epoch = checkpoint['epoch'] + del optimizer, scheduler + optimizer, scheduler = fetch_optimizer(args, model, start_step, checkpoint) + + train_data = datasets.fetch_dataset(args) + if args.distributed: + train_sampler = torch.utils.data.distributed.DistributedSampler( + train_data, + num_replicas=torch.cuda.device_count(), + rank=args.local_rank + ) + else: + train_sampler = None + train_loader = DataLoader(dataset=train_data, batch_size=args.batch_size, shuffle=train_sampler is None, + num_workers=args.num_workers, pin_memory=True, drop_last=True, + sampler=train_sampler) + + total_steps = start_step + epoch = start_epoch + logger = Logger(model, scheduler, args.name) + logger.total_steps = total_steps + + model.train() + scaler = GradScaler(enabled=args.mixed_precision) + should_keep_training = True + + while should_keep_training: + + # mannually change random seed for shuffling every epoch + if args.distributed: + train_sampler.set_epoch(epoch) + + if total_steps == start_step: + epoch_start_step = start_step - len(train_loader)*start_epoch + else: + epoch_start_step = 0 + + for i_batch, data_blob in enumerate(tqdm(train_loader, initial=epoch_start_step)): + optimizer.zero_grad() + image1 = data_blob["img1"].cuda() + image2 = data_blob["img2"].cuda() + disp_gt = data_blob["disp"].cuda() + valid = data_blob["valid"].cuda() + + assert model.training + disp_predictions = model(image1, image2, iters=args.train_iters, scale_iters=args.scale_iters) + assert model.training + + loss, metrics = sequence_loss(disp_predictions, disp_gt, valid) + + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + + scaler.step(optimizer) + scheduler.step() + scaler.update() + + total_steps += 1 + + if args.local_rank == 0: + logger.writer.add_scalar("train/live_loss", loss.item(), total_steps) + logger.writer.add_scalar(f'train/learning_rate', optimizer.param_groups[0]['lr'], total_steps) + logger.push(metrics) + + if total_steps % args.save_latest_ckpt_freq == 0: + save_path = Path('checkpoints/%s/checkpoint_latest.pth' % (args.name)) + logging.info(f"Saving file {save_path.absolute()}") + save_dict = { 'model': model_without_ddp.state_dict(), + 'optimizer': optimizer.state_dict(), + 'step': total_steps, + 'epoch': epoch} + torch.save(save_dict, save_path) + + if total_steps % args.save_ckpt_freq == 0: + save_path = Path('checkpoints/%s/%s_%6d.pth' % (args.name, args.name, total_steps)) + logging.info(f"Saving file {save_path.absolute()}") + torch.save(model_without_ddp.state_dict(), save_path) + + if total_steps % args.val_freq == 0: + + # visualizing training results with tensorboard + disp = disp_predictions[-1] + + for j in range(min(4, args.batch_size)): # write a maxmimum of four images + logger.writer.add_image("image1/{}".format(j), image1[j].data.type(torch.uint8), total_steps) + logger.writer.add_image("image2/{}".format(j), image2[j].data.type(torch.uint8), total_steps) + logger.writer.add_image("disp/{}".format(j), + (disp[j]).data.type(torch.uint8), total_steps) + logger.writer.add_image("gt_disp/{}".format(j), + (disp_gt[j]).data.type(torch.uint8), total_steps) + + results = validate_things(model_without_ddp, args.valid_iters, args.scale_iters) + logger.write_dict(results) + model.train() + if not args.distributed: model_without_ddp.freeze_bn() + + if total_steps > args.num_steps: + should_keep_training = False + break + + epoch += 1 + + if len(train_loader) >= 10000: + save_path = Path('checkpoints/%s/%d_epoch_%s.pth.gz' % (args.name, total_steps, args.name)) + logging.info(f"Saving file {save_path}") + torch.save(model_without_ddp.state_dict(), save_path) + + print("FINISHED TRAINING") + logger.close() + PATH = 'checkpoints/%s.pth' % args.name + torch.save(model_without_ddp.state_dict(), PATH) + + return PATH + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--name', default='defom-stereo', help="name your experiment") + + # resume pretrained model or resume training + parser.add_argument('--resume_ckpt', default=None, type=str, + help='resume from pretrained model or resume from unexpectedly terminated training') + parser.add_argument('--strict_resume', action='store_true', + help='strict resume while loading pretrained weights') + parser.add_argument('--no_resume_optimizer', action='store_true') + + # Training parameters + parser.add_argument('--batch_size', type=int, default=8, help="batch size used during training.") + parser.add_argument('--num_workers', default=8, type=int) + parser.add_argument('--train_datasets', nargs='+', default=['sceneflow'], help="training datasets.") + parser.add_argument('--train_folds', type=int, nargs='+', default=[1], help="training datasets' folds.") + parser.add_argument('--lr', type=float, default=0.0002, help="max learning rate.") + parser.add_argument('--image_size', type=int, nargs='+', default=[320, 736], help="size of the random image crops used during training.") + parser.add_argument('--train_iters', type=int, default=18, help="number of updates to the disparity field in each forward pass.") + parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.") + parser.add_argument('--wdecay', type=float, default=.00001, help="Weight decay in optimizer.") + parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') + parser.add_argument('--seed', default=1234, type=int) + + # log + parser.add_argument('--num_steps', type=int, default=200000, help="length of training schedule.") + parser.add_argument('--save_ckpt_freq', default=10000, type=int, help='Save checkpoint frequency (steps)') + parser.add_argument('--save_latest_ckpt_freq', default=1000, type=int) + parser.add_argument('--val_freq', default=10000, type=int, help='validation frequency in terms of training steps') + + # distributed training + parser.add_argument('--distributed', action='store_true') + parser.add_argument('--local-rank', type=int, default=0) + parser.add_argument('--launcher', default='none', type=str) + parser.add_argument('--gpu_ids', default=0, type=int, nargs='+') + + # Validation parameters + parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during validation forward pass') + + # Raft Architecure choices + parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg']) + parser.add_argument('--idepth_scale', type=float, default=0.5, help="the scale of inverse depth to initialize disparity") + parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation") + parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid") + parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid") + + parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0], + help='the list of scaling factors of disparity') + parser.add_argument('--scale_corr_radius', type=int, default=2, help="width of the correlation pyramid for scaled disparity") + + parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], help="resolution of the disparity field (1/2^K)") + parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder") + parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels") + parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions") + + # Data augmentation + parser.add_argument('--img_gamma', type=float, nargs='+', default=None, help="gamma range") + parser.add_argument('--saturation_range', type=float, nargs='+', default=[0.0, 1.4], help='color saturation') + parser.add_argument('--do_flip', default='v', choices=['v', 'None'], help='flip the images vertically') + parser.add_argument('--spatial_scale', type=float, nargs='+', default=[-0.2, 0.4], help='re-scale the images randomly') + parser.add_argument('--noyjitter', action='store_true', help='don\'t simulate imperfect rectification') + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO, + stream=sys.stdout, + format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s') + + if 'LOCAL_RANK' not in os.environ: + os.environ['LOCAL_RANK'] = str(args.local_rank) + + Path("checkpoints/"+args.name).mkdir(exist_ok=True, parents=True) + + train(args) diff --git a/cv/defom-stereo/utils/dist_utils.py b/cv/defom-stereo/utils/dist_utils.py new file mode 100644 index 0000000..fdb08c7 --- /dev/null +++ b/cv/defom-stereo/utils/dist_utils.py @@ -0,0 +1,105 @@ +# Copyright (c) OpenMMLab. All rights reserved. +# https://github.com/open-mmlab/mmcv/blob/7540cf73ac7e5d1e14d0ffbd9b6759e83929ecfc/mmcv/runner/dist_utils.py + +import os +import subprocess + +import torch +import torch.multiprocessing as mp +from torch import distributed as dist + + +def init_dist(launcher, backend='nccl', **kwargs): + if mp.get_start_method(allow_none=True) is None: + mp.set_start_method('spawn') + if launcher == 'pytorch': + _init_dist_pytorch(backend, **kwargs) + elif launcher == 'mpi': + _init_dist_mpi(backend, **kwargs) + elif launcher == 'slurm': + _init_dist_slurm(backend, **kwargs) + else: + raise ValueError(f'Invalid launcher type: {launcher}') + + +def _init_dist_pytorch(backend, **kwargs): + # TODO: use local_rank instead of rank % num_gpus + rank = int(os.environ['RANK']) + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(rank % num_gpus) + dist.init_process_group(backend=backend, **kwargs) + + +def _init_dist_mpi(backend, **kwargs): + # TODO: use local_rank instead of rank % num_gpus + rank = int(os.environ['OMPI_COMM_WORLD_RANK']) + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(rank % num_gpus) + dist.init_process_group(backend=backend, **kwargs) + + +def _init_dist_slurm(backend, port=None): + """Initialize slurm distributed training environment. + If argument ``port`` is not specified, then the master port will be system + environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system + environment variable, then a default port ``29500`` will be used. + Args: + backend (str): Backend of torch.distributed. + port (int, optional): Master port. Defaults to None. + """ + proc_id = int(os.environ['SLURM_PROCID']) + ntasks = int(os.environ['SLURM_NTASKS']) + node_list = os.environ['SLURM_NODELIST'] + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(proc_id % num_gpus) + addr = subprocess.getoutput( + f'scontrol show hostname {node_list} | head -n1') + # specify master port + if port is not None: + os.environ['MASTER_PORT'] = str(port) + elif 'MASTER_PORT' in os.environ: + pass # use MASTER_PORT in the environment variable + else: + # 29500 is torch.distributed default port + os.environ['MASTER_PORT'] = '29500' + # use MASTER_ADDR in the environment variable if it already exists + if 'MASTER_ADDR' not in os.environ: + os.environ['MASTER_ADDR'] = addr + os.environ['WORLD_SIZE'] = str(ntasks) + os.environ['LOCAL_RANK'] = str(proc_id % num_gpus) + os.environ['RANK'] = str(proc_id) + dist.init_process_group(backend=backend) + + +def get_dist_info(): + # if (TORCH_VERSION != 'parrots' + # and digit_version(TORCH_VERSION) < digit_version('1.0')): + # initialized = dist._initialized + # else: + if dist.is_available(): + initialized = dist.is_initialized() + else: + initialized = False + if initialized: + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + return rank, world_size + + +# from DETR repo +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop('force', False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print diff --git a/cv/defom-stereo/utils/utils.py b/cv/defom-stereo/utils/utils.py new file mode 100644 index 0000000..56cc656 --- /dev/null +++ b/cv/defom-stereo/utils/utils.py @@ -0,0 +1,119 @@ +import logging +import sys + +import torch +import torch.optim as optim +from torch.utils.tensorboard import SummaryWriter + +import numpy as np +import random + +def seed_everything(seed): + torch.manual_seed(seed) # Current CPU + torch.cuda.manual_seed(seed) # Current GPU + np.random.seed(seed) # Numpy module + random.seed(seed) # Python random module + torch.backends.cudnn.benchmark = False # Close optimization + torch.backends.cudnn.deterministic = True # Close optimization + torch.cuda.manual_seed_all(seed) # All GPU (Optional) + + +def sequence_loss(flow_preds, flow_gt, valid, loss_gamma=0.9, max_flow=700): + """ Loss function defined over sequence of flow predictions """ + + n_predictions = len(flow_preds) + assert n_predictions >= 1 + flow_loss = 0.0 + + # exlude invalid pixels and extremely large diplacements + mag = torch.sum(flow_gt ** 2, dim=1, keepdim=True).sqrt() + + # exclude extremly large displacements + valid = ((valid >= 0.5) & (mag < max_flow)) + assert valid.shape == flow_gt.shape, [valid.shape, flow_gt.shape] + assert not torch.isinf(flow_gt[valid.bool()]).any() + + for i in range(n_predictions): + assert not torch.isnan(flow_preds[i]).any() and not torch.isinf(flow_preds[i]).any() + # We adjust the loss_gamma so it is consistent for any number of RAFT-Stereo iterations + adjusted_loss_gamma = loss_gamma ** (15 / (n_predictions)) + i_weight = adjusted_loss_gamma ** (n_predictions - i) + i_loss = (flow_preds[i] - flow_gt).abs() + assert i_loss.shape == valid.shape, [i_loss.shape, valid.shape, flow_gt.shape, flow_preds[i].shape] + flow_loss += i_weight * i_loss[valid.bool()].mean() + + epe = torch.sum((flow_preds[-1] - flow_gt) ** 2, dim=1).sqrt() + epe = epe.view(-1)[valid.view(-1)] + + metrics = { + 'epe': epe.mean().item(), + '1px': (epe < 1).float().mean().item(), + '3px': (epe < 3).float().mean().item(), + '5px': (epe < 5).float().mean().item(), + } + + return flow_loss, metrics + + +def fetch_optimizer(args, model, last_epoch=-1, checkpoint=None): + """ Create the optimizer and learning rate scheduler """ + trainable_params = filter(lambda p: p.requires_grad, model.parameters()) + optimizer = optim.AdamW(trainable_params, lr=args.lr, weight_decay=args.wdecay, eps=1e-8) + if checkpoint is not None: + optimizer.load_state_dict(checkpoint['optimizer']) + + scheduler = optim.lr_scheduler.OneCycleLR(optimizer, args.lr, args.num_steps + 100, pct_start=0.01, + cycle_momentum=False, anneal_strategy='linear', last_epoch=last_epoch) + + return optimizer, scheduler + + +class Logger: + SUM_FREQ = 100 + + def __init__(self, model, scheduler, name): + self.model = model + self.scheduler = scheduler + self.total_steps = 0 + self.running_loss = {} + self.log_dir = 'runs/' + name + self.writer = SummaryWriter(log_dir=self.log_dir) + + def _print_training_status(self): + metrics_data = [self.running_loss[k] / Logger.SUM_FREQ for k in sorted(self.running_loss.keys())] + training_str = "[{:6d}, {:10.7f}] ".format(self.total_steps + 1, self.scheduler.get_last_lr()[0]) + metrics_str = ("{:10.4f}, " * len(metrics_data)).format(*metrics_data) + + # print the training status + logging.info(f"Training Metrics ({self.total_steps}): {training_str + metrics_str}") + + if self.writer is None: + self.writer = SummaryWriter(log_dir=self.log_dir) + + for k in self.running_loss: + self.writer.add_scalar("train/" + k, self.running_loss[k] / Logger.SUM_FREQ, self.total_steps) + self.running_loss[k] = 0.0 + + def push(self, metrics): + self.total_steps += 1 + + for key in metrics: + if key not in self.running_loss: + self.running_loss[key] = 0.0 + + self.running_loss[key] += metrics[key] + + if self.total_steps % Logger.SUM_FREQ == Logger.SUM_FREQ - 1: + self._print_training_status() + self.running_loss = {} + + def write_dict(self, results): + if self.writer is None: + self.writer = SummaryWriter(log_dir=self.log_dir) + + for key in results: + self.writer.add_scalar("valid/" + key, results[key], self.total_steps) + + def close(self): + self.writer.close() + diff --git a/cv/fast-foundationstereo/.gitignore b/cv/fast-foundationstereo/.gitignore new file mode 100644 index 0000000..7420e2f --- /dev/null +++ b/cv/fast-foundationstereo/.gitignore @@ -0,0 +1,14 @@ +__pycache__ +*.vscode +cfg.yaml +*.pth +*.pkl +*.pyc +*.so +*.egg-info +# Ignore weight files but keep the folder structure +weights/* +!weights/.gitkeep +output/ +.claude/ +cpp/build/ \ No newline at end of file diff --git a/cv/fast-foundationstereo/LICENSE.txt b/cv/fast-foundationstereo/LICENSE.txt new file mode 100644 index 0000000..db16c96 --- /dev/null +++ b/cv/fast-foundationstereo/LICENSE.txt @@ -0,0 +1,94 @@ +Copyright (c) 2026-Present, NVIDIA Corporation & affiliates. All rights reserved. + + +======================================================================= + +1. Definitions + +"Licensor" means any person or entity that distributes its Work. + +"Software" means the original work of authorship made available under +this License. + +"Work" means the Software and any additions to or derivative works of +the Software that are made available under this License. + +The terms "reproduce," "reproduction," "derivative works," and +"distribution" have the meaning as provided under U.S. copyright law; +provided, however, that for the purposes of this License, derivative +works shall not include works that remain separable from, or merely +link (or bind by name) to the interfaces of, the Work. + +Works, including the Software, are "made available" under this License +by including in or with the Work either (a) a copyright notice +referencing the applicability of this License to the Work, or (b) a +copy of this License. + +2. License Grants + + 2.1 Copyright Grant. Subject to the terms and conditions of this + License, each Licensor grants to you a perpetual, worldwide, + non-exclusive, royalty-free, copyright license to reproduce, + prepare derivative works of, publicly display, publicly perform, + sublicense and distribute its Work and any resulting derivative + works in any form. + +3. Limitations + + 3.1 Redistribution. You may reproduce or distribute the Work only + if (a) you do so under this License, (b) you include a complete + copy of this License with your distribution, and (c) you retain + without modification any copyright, patent, trademark, or + attribution notices that are present in the Work. + + 3.2 Derivative Works. You may specify that additional or different + terms apply to the use, reproduction, and distribution of your + derivative works of the Work ("Your Terms") only if (a) Your Terms + provide that the use limitation in Section 3.3 applies to your + derivative works, and (b) you identify the specific derivative + works that are subject to Your Terms. Notwithstanding Your Terms, + this License (including the redistribution requirements in Section + 3.1) will continue to apply to the Work itself. + + 3.3 Use Limitation. The Work and any derivative works thereof only + may be used or intended for use non-commercially. Notwithstanding + the foregoing, NVIDIA and its affiliates may use the Work and any + derivative works commercially. As used herein, "non-commercially" + means for research purposes only. + + 3.4 Patent Claims. If you bring or threaten to bring a patent claim + against any Licensor (including any claim, cross-claim or + counterclaim in a lawsuit) to enforce any patents that you allege + are infringed by any Work, then your rights under this License from + such Licensor (including the grant in Section 2.1) will terminate + immediately. + + 3.5 Trademarks. This License does not grant any rights to use any + Licensor�s or its affiliates� names, logos, or trademarks, except + as necessary to reproduce the notices described in this License. + + 3.6 Termination. If you violate any term of this License, then your + rights under this License (including the grant in Section 2.1) will + terminate immediately. + +4. Disclaimer of Warranty. + +THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR +NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER +THIS LICENSE. + +5. Limitation of Liability. + +EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL +THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE +SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, +INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF +OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK +(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, +LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER +COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF +THE POSSIBILITY OF SUCH DAMAGES. + +======================================================================= diff --git a/cv/fast-foundationstereo/Utils.py b/cv/fast-foundationstereo/Utils.py new file mode 100755 index 0000000..1f12387 --- /dev/null +++ b/cv/fast-foundationstereo/Utils.py @@ -0,0 +1,88 @@ +import os, sys, torch, imageio, logging, importlib, argparse +import cv2 +import numpy as np +import yaml +try: + import open3d as o3d +except: + o3d = None + +AMP_DTYPE = torch.float16 + +def set_logging_format(level=logging.INFO): + importlib.reload(logging) + FORMAT = '%(message)s' + logging.basicConfig(level=level, format=FORMAT, datefmt='%m-%d|%H:%M:%S') + + +def set_seed(random_seed): + import torch,random + np.random.seed(random_seed) + random.seed(random_seed) + torch.manual_seed(random_seed) + torch.cuda.manual_seed_all(random_seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def toOpen3dCloud(points,colors=None,normals=None): + cloud = o3d.geometry.PointCloud() + cloud.points = o3d.utility.Vector3dVector(points.astype(np.float64)) + if colors is not None: + if colors.max()>1: + colors = colors/255.0 + cloud.colors = o3d.utility.Vector3dVector(colors.astype(np.float64)) + if normals is not None: + cloud.normals = o3d.utility.Vector3dVector(normals.astype(np.float64)) + return cloud + + + +def depth2xyzmap(depth:np.ndarray, K, uvs:np.ndarray=None, zmin=0.1): + invalid_mask = (depth thres is invalid + """ + disp = disp.copy() + H,W = disp.shape[:2] + invalid_mask = disp>=invalid_thres + if (invalid_mask==0).sum()==0: + other_output['min_val'] = None + other_output['max_val'] = None + return np.zeros((H,W,3)) + if min_val is None: + min_val = disp[invalid_mask==0].min() + if max_val is None: + max_val = disp[invalid_mask==0].max() + other_output['min_val'] = min_val + other_output['max_val'] = max_val + vis = ((disp-min_val)/(max_val-min_val)).clip(0,1) * 255 + if cmap is None: + vis = cv2.applyColorMap(vis.clip(0, 255).astype(np.uint8), color_map)[...,::-1] + else: + vis = cmap(vis.astype(np.uint8))[...,:3]*255 + if invalid_mask.any(): + vis[invalid_mask] = 0 + return vis.astype(np.uint8) diff --git a/cv/fast-foundationstereo/assets/bp2_vs_runtime.jpg b/cv/fast-foundationstereo/assets/bp2_vs_runtime.jpg new file mode 100644 index 0000000..e17e6b4 Binary files /dev/null and b/cv/fast-foundationstereo/assets/bp2_vs_runtime.jpg differ diff --git a/cv/fast-foundationstereo/assets/disp_vis.png b/cv/fast-foundationstereo/assets/disp_vis.png new file mode 100755 index 0000000..138dd1a Binary files /dev/null and b/cv/fast-foundationstereo/assets/disp_vis.png differ diff --git a/cv/fast-foundationstereo/assets/intro.jpg b/cv/fast-foundationstereo/assets/intro.jpg new file mode 100644 index 0000000..4b813d4 Binary files /dev/null and b/cv/fast-foundationstereo/assets/intro.jpg differ diff --git a/cv/fast-foundationstereo/assets/intro_c.mp4 b/cv/fast-foundationstereo/assets/intro_c.mp4 new file mode 100644 index 0000000..e7d1a99 Binary files /dev/null and b/cv/fast-foundationstereo/assets/intro_c.mp4 differ diff --git a/cv/fast-foundationstereo/assets/intro_c.webp b/cv/fast-foundationstereo/assets/intro_c.webp new file mode 100644 index 0000000..de61cba Binary files /dev/null and b/cv/fast-foundationstereo/assets/intro_c.webp differ diff --git a/cv/fast-foundationstereo/assets/pcl_vis.png b/cv/fast-foundationstereo/assets/pcl_vis.png new file mode 100755 index 0000000..949386c Binary files /dev/null and b/cv/fast-foundationstereo/assets/pcl_vis.png differ diff --git a/cv/fast-foundationstereo/assets/stereo4d.jpg b/cv/fast-foundationstereo/assets/stereo4d.jpg new file mode 100644 index 0000000..f99e415 Binary files /dev/null and b/cv/fast-foundationstereo/assets/stereo4d.jpg differ diff --git a/cv/fast-foundationstereo/assets/stereo4d_labeling.jpg b/cv/fast-foundationstereo/assets/stereo4d_labeling.jpg new file mode 100644 index 0000000..c39a347 Binary files /dev/null and b/cv/fast-foundationstereo/assets/stereo4d_labeling.jpg differ diff --git a/cv/fast-foundationstereo/core/__init__.py b/cv/fast-foundationstereo/core/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/cv/fast-foundationstereo/core/distill_block.py b/cv/fast-foundationstereo/core/distill_block.py new file mode 100755 index 0000000..dc2dd52 --- /dev/null +++ b/cv/fast-foundationstereo/core/distill_block.py @@ -0,0 +1,50 @@ +import os,sys +code_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(code_dir+'/../') +from foundation_stereo_ori.submodule import FeatureAtt +import torch +import torch.nn as nn +import Utils as U +import pickle + + +class ForwardHelper(nn.Module): + def __init__(self, layers:list): + super().__init__() + self.layers = nn.ModuleList(layers) + + def forward(self, x, left_feat=None): + for layer in self.layers: + if isinstance(layer, FeatureAtt): + x = layer(x, left_feat) + else: + x = layer(x) + return x + + +class PostForwardHelper(nn.Module): + def __init__(self, layers:list): + super().__init__() + for pos in range(len(layers)): + if layers[pos] in ['sum', 'concat']: + self.op = layers[pos] + break + self.upsample = nn.Sequential(*layers[:pos]) + self.out = nn.ModuleList(layers[pos+1:]) + + def forward(self, conv2, conv3, left_feat=None): + conv3_up = self.upsample(conv3) + if self.op == 'sum': + x = conv3_up + conv2 + elif self.op == 'concat': + x = torch.cat((conv3_up, conv2), dim=1) + else: + raise ValueError(f"Unknown operation: {self.op}") + + for layer in self.out: + if isinstance(layer, FeatureAtt): + x = layer(x, left_feat) + else: + x = layer(x) + return x + diff --git a/cv/fast-foundationstereo/core/extractor.py b/cv/fast-foundationstereo/core/extractor.py new file mode 100755 index 0000000..39d928a --- /dev/null +++ b/cv/fast-foundationstereo/core/extractor.py @@ -0,0 +1,78 @@ +import torch,os,sys +import torch.nn as nn +import torch.nn.functional as F +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from core.submodule import Conv2x_IN +import timm + + + +class ContextNetSharedBackbone(nn.Module): + def __init__(self, args, c04, c08, c16, output_dim=[(128,128,128), (128,128,128)], norm_fn='batch', downsample=3): + super().__init__() + self.args = args + self.conv04 = nn.ModuleList([ + nn.Conv2d(c04, output_dim[0][0], kernel_size=3, padding=1), + nn.Conv2d(c04, output_dim[1][0], kernel_size=3, padding=1), + ]) + + def forward(self, x4, x8, x16): + outputs04 = [] + for i in range(len(self.conv04)): + outputs04.append(self.conv04[i](x4)) + return (outputs04,) + + + +class DepthAnythingFeature: + model_configs = { + 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}, + 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, + 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]} + } + + + +class Feature(nn.Module): + def __init__(self, args): + super(Feature, self).__init__() + self.args = args + model = timm.create_model('edgenext_small', pretrained=True, features_only=False) + self.stem = model.stem + self.stages = model.stages + chans = [48, 96, 160, 304] + self.chans = chans + vit_feat_dim = DepthAnythingFeature.model_configs[self.args.vit_size]['features']//2 + + self.deconv32_16 = Conv2x_IN(chans[3], chans[2], deconv=True, concat=True) + self.deconv16_8 = Conv2x_IN(chans[2]*2, chans[1], deconv=True, concat=True) + self.deconv8_4 = Conv2x_IN(chans[1]*2, chans[0], deconv=True, concat=True) + + self.conv4 = nn.Conv2d(chans[0]*2, self.chans[0]*2+vit_feat_dim, kernel_size=1, stride=1, padding=0) + + self.d_out = [self.chans[0]*2+vit_feat_dim, self.chans[1]*2, self.chans[2]*2, self.chans[3]] + + + def forward(self, x): + B,C,H,W = x.shape + if hasattr(self, 'stem'): + x = self.stem(x) + x4 = self.stages[0](x) + x8 = self.stages[1](x4) + x16 = self.stages[2](x8) + x32 = self.stages[3](x16) + else: + intermediates = self.model.forward_intermediates(x, intermediates_only=True) + x4, x8, x16, x32 = intermediates[-4:] + + with torch.profiler.record_function("feature_deconv"): + x16 = self.deconv32_16(x32, x16) + x8 = self.deconv16_8(x16, x8) + x4 = self.deconv8_4(x8, x4) + x4 = self.conv4(x4) + if hasattr(self, 'conv8'): + x8 = self.conv8(x8) + x16 = self.conv16(x16) + x32 = self.conv32(x32) + return [x4, x8, x16, x32] diff --git a/cv/fast-foundationstereo/core/foundation_stereo.py b/cv/fast-foundationstereo/core/foundation_stereo.py new file mode 100755 index 0000000..b1377d4 --- /dev/null +++ b/cv/fast-foundationstereo/core/foundation_stereo.py @@ -0,0 +1,445 @@ +import torch,pdb,logging,timm +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import sys,os +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from core.update import BasicSelectiveMultiUpdateBlock +from core.extractor import ContextNetSharedBackbone, Feature +from core.geometry import Combined_Geo_Encoding_Volume +from core.submodule import ( + BasicConv, Conv3dNormActReduced, ResnetBasicBlock3D, BasicConv_IN, Conv2x, + FeatureAtt, CostVolumeDisparityAttention, SpatialAttentionExtractor, + ChannelAttentionEnhancement, disparity_regression, context_upsample, + build_gwc_volume_optimized_pytorch1, build_gwc_volume_triton, + build_concat_volume_optimized_pytorch1, build_concat_volume_optimized_pytorch, +) +from core.utils.utils import InputPadder +import Utils as U +import time + +sys.modules['foundation_stereo_ori'] = sys.modules['core'] +sys.modules['foundation_stereo_ori.submodule'] = sys.modules['core.submodule'] +sys.modules['foundation_stereo_ori.extractor'] = sys.modules['core.extractor'] +sys.modules['foundation_stereo_ori.update'] = sys.modules['core.update'] +sys.modules['foundation_stereo_ori.foundation_stereo'] = sys.modules['core.foundation_stereo'] +class FoundationStereo(nn.Module): + pass + + + +def normalize_image(img): + ''' + @img: (B,C,H,W) in range 0-255, RGB order + ''' + mean = img.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = img.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + return (img/255.0 - mean) / std + + +class hourglass(nn.Module): + def __init__(self, cfg, in_channels, feat_dims=None): + super().__init__() + self.cfg = cfg + self.conv1 = nn.Sequential(BasicConv(in_channels, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17)) + + self.conv2 = nn.Sequential(BasicConv(in_channels*2, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17)) + + self.conv3 = nn.Sequential(BasicConv(in_channels*4, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + Conv3dNormActReduced(in_channels*6, in_channels*6, kernel_size=3, kernel_disp=17)) + + + self.conv3_up = BasicConv(in_channels*6, in_channels*4, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.conv2_up = BasicConv(in_channels*4, in_channels*2, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.conv1_up = BasicConv(in_channels*2, in_channels, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + self.conv_out = nn.Sequential( + Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17), + Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17), + ) + + self.agg_0 = nn.Sequential(BasicConv(in_channels*8, in_channels*4, is_3d=True, kernel_size=1, padding=0, stride=1), + Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17), + Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),) + + self.agg_1 = nn.Sequential(BasicConv(in_channels*4, in_channels*2, is_3d=True, kernel_size=1, padding=0, stride=1), + Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17), + Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17)) + + self.atts = nn.ModuleDict({ + "4": CostVolumeDisparityAttention(d_model=in_channels, nhead=4, dim_feedforward=in_channels, norm_first=False, num_transformer=4, max_len=self.cfg['max_disp']//16), + }) + self.conv_patch = nn.Sequential( + nn.Conv3d(in_channels, in_channels, kernel_size=4, stride=4, padding=0, groups=in_channels), + nn.BatchNorm3d(in_channels), + ) + self.feature_att_8 = FeatureAtt(in_channels*2, feat_dims[1]) + self.feature_att_16 = FeatureAtt(in_channels*4, feat_dims[2]) + self.feature_att_32 = FeatureAtt(in_channels*6, feat_dims[3]) + self.feature_att_up_16 = FeatureAtt(in_channels*4, feat_dims[2]) + self.feature_att_up_8 = FeatureAtt(in_channels*2, feat_dims[1]) + + self.post32_to_16 = None + self.post16_to_8 = None + self.post8_to_4 = None + + def forward(self, x, features): + conv1 = self.conv1(x) + conv1 = self.feature_att_8(conv1, features[1]) + + conv2 = self.conv2(conv1) + conv2 = self.feature_att_16(conv2, features[2]) + + conv3 = self.conv3(conv2) + conv3 = self.feature_att_32(conv3, features[3]) + if not hasattr(self, 'post32_to_16') or self.post32_to_16 is None: + conv3_up = self.conv3_up(conv3) + conv2 = torch.cat((conv3_up, conv2), dim=1) + conv2 = self.agg_0(conv2) + conv2 = self.feature_att_up_16(conv2, features[2]) + else: + conv2 = self.post32_to_16(conv2, conv3, features[2]) + + if not hasattr(self, 'post16_to_8') or self.post16_to_8 is None: + conv2_up = self.conv2_up(conv2) + conv1 = torch.cat((conv2_up, conv1), dim=1) + conv1 = self.agg_1(conv1) + conv1 = self.feature_att_up_8(conv1, features[1]) + else: + conv1 = self.post16_to_8(conv1, conv2, features[1]) + + conv = self.conv1_up(conv1) + if not hasattr(self, 'post8_to_4') or self.post8_to_4 is None: + x = self.conv_patch(x) + x = self.atts["4"](x) + x = F.interpolate(x, scale_factor=4, mode='trilinear', align_corners=False) + conv = conv + x + conv = self.conv_out(conv) + else: + conv = self.post8_to_4(x, conv) + + return conv + + +class FastFoundationStereo(nn.Module): + def __init__(self, args): + super().__init__() + self.args = args + self.dtype = torch.float32 + + context_dims = args.hidden_dims + self.cv_group = args.get('cv_group', 8) + self.concat_channel = 24 + volume_dim = args.get('volume_dim', 28) + self.volume_dim = volume_dim + self.update_block = BasicSelectiveMultiUpdateBlock(self.args, self.args.hidden_dims[0], volume_dim=volume_dim) + self.sam = SpatialAttentionExtractor() + self.cam = ChannelAttentionEnhancement(self.args.hidden_dims[0]) + self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, kernel_size=3, padding=3//2) for i in range(self.args.n_gru_layers)]) + self.feature = Feature(args) + self.proj_cmb = nn.Conv2d(self.feature.d_out[0], self.concat_channel//2, kernel_size=1, padding=0) + self.cnet = ContextNetSharedBackbone(args, c04=self.feature.d_out[0], c08=self.feature.d_out[1], c16=self.feature.d_out[2], output_dim=[args.hidden_dims, context_dims]) + + self.stem_2 = nn.Sequential( + BasicConv_IN(3, 32, kernel_size=3, stride=2, padding=1), + nn.Conv2d(32, 32, 3, 1, 1, bias=False), + nn.InstanceNorm2d(32), nn.ReLU() + ) + self.spx_2_gru = Conv2x(32, 32, deconv=True, bn=False, concat=True) + self.spx_gru = nn.Sequential( + nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1), + ) + + self.corr_stem = nn.Sequential( + nn.Conv3d(self.proj_cmb.out_channels*2+self.cv_group, volume_dim, kernel_size=1), + BasicConv(volume_dim, volume_dim, kernel_size=3, padding=1, is_3d=True), + ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1), + ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1), + ) + self.corr_feature_att = FeatureAtt(volume_dim, self.feature.d_out[0]) + self.cost_agg = hourglass(cfg=self.args, in_channels=volume_dim, feat_dims=self.feature.d_out) + self.classifier = nn.Sequential( + BasicConv(volume_dim, volume_dim//2, kernel_size=3, padding=1, is_3d=True), + ResnetBasicBlock3D(volume_dim//2, volume_dim//2, kernel_size=3, stride=1, padding=1), + nn.Conv3d(volume_dim//2, 1, kernel_size=7, padding=3), + ) + + r = self.args.corr_radius + dx = torch.arange(-r, r+1, requires_grad=False, dtype=torch.int8).reshape(1, 1, 2*r+1, 1) + self.register_buffer("dx", dx) + + + def upsample_disp(self, disp, mask_feat_4, stem_2x): + with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE): + xspx = self.spx_2_gru(mask_feat_4, stem_2x) # 1/2 resolution + spx_pred = self.spx_gru(xspx) + spx_pred = F.softmax(spx_pred, 1) + up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1) + return up_disp.to(self.dtype) + + + def forward(self, image1, image2, iters=12, test_mode=False, low_memory=False, init_disp=None, profile=False, optimize_build_volume='pytorch1'): + """ Estimate disparity between pair of frames """ + B,C,H,W = image1.shape + low_memory = low_memory or (self.args.get('low_memory', False)) + image1 = normalize_image(image1) + image2 = normalize_image(image2) + with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE): + out = self.feature(torch.cat([image1, image2], dim=0)) + features_left = [o[:B] for o in out] + features_right = [o[B:] for o in out] + stem_2x = self.stem_2(image1) + + if optimize_build_volume=='pytorch1': + gwc_volume = build_gwc_volume_optimized_pytorch1(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group, normalize=self.args.normalize) + elif optimize_build_volume=='triton': + gwc_volume = build_gwc_volume_triton(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group, normalize=self.args.normalize) + else: + raise RuntimeError(f"Invalid optimize_build_volume: {optimize_build_volume}") + + left_tmp = self.proj_cmb(features_left[0]) + right_tmp = self.proj_cmb(features_right[0]) + concat_volume = build_concat_volume_optimized_pytorch1(left_tmp, right_tmp, maxdisp=self.args.max_disp//4) + del left_tmp, right_tmp + comb_volume = torch.cat([gwc_volume, concat_volume], dim=1) + del concat_volume, gwc_volume + + comb_volume = self.corr_stem(comb_volume) + comb_volume = self.corr_feature_att(comb_volume, features_left[0]) + comb_volume = self.cost_agg(comb_volume, features_left) + + # Init disp from geometry encoding volume + logits = self.classifier(comb_volume).squeeze(1) + prob = F.softmax(logits, dim=1) + if init_disp is None: + init_disp = disparity_regression(prob, self.args.max_disp//4) + + cnet_list = self.cnet(features_left[0], features_left[1], features_left[2]) + cnet_list = list(cnet_list) + net_list = [torch.tanh(x[0]) for x in cnet_list] + inp_list = [torch.relu(x[1]) for x in cnet_list] + inp_list = [self.cam(x) * x for x in inp_list] + att = [self.sam(x) for x in inp_list] + + geo_fn = Combined_Geo_Encoding_Volume(features_left[0].to(self.dtype), features_right[0].to(self.dtype), comb_volume.to(self.dtype), num_levels=self.args.corr_levels) + b, c, h, w = features_left[0].shape + coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1) + disp = init_disp.to(self.dtype) + disp_preds = [] + + del comb_volume, features_left, features_right, cnet_list + + # GRUs iterations to update disparity (1/4 resolution) + for itr in range(iters): + disp = disp.detach() + geo_feat = geo_fn(disp, coords, dx=self.dx, low_memory=low_memory) + with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE): + net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat.to(self.dtype), disp, att) + + disp = disp + delta_disp.to(self.dtype) + if test_mode and itr < iters-1: + continue + + # upsample predictions + disp_up = self.upsample_disp(disp.to(self.dtype), mask_feat_4.to(self.dtype), stem_2x.to(self.dtype)) + disp_preds.append(disp_up) + + + if test_mode: + return disp_up + + return init_disp, disp_preds + + + def run_hierachical(self, image1, image2, iters=12, test_mode=False, low_memory=False, small_ratio=0.5): + B,_,H,W = image1.shape + img1_small = F.interpolate(image1, scale_factor=small_ratio, align_corners=False, mode='bilinear') + img2_small = F.interpolate(image2, scale_factor=small_ratio, align_corners=False, mode='bilinear') + padder = InputPadder(img1_small.shape[-2:], divis_by=32, force_square=False) + img1_small, img2_small = padder.pad(img1_small, img2_small) + disp_small = self.forward(img1_small, img2_small, test_mode=True, iters=iters, low_memory=low_memory) + disp_small = padder.unpad(disp_small) + disp_small_up = F.interpolate(disp_small, size=(H,W), mode='bilinear', align_corners=True) * 1/small_ratio + disp_small_up = disp_small_up.clip(0, None) + + padder = InputPadder(image1.shape[-2:], divis_by=32, force_square=False) + image1, image2, disp_small_up = padder.pad(image1, image2, disp_small_up) + disp_small_up += padder._pad[0] + init_disp = F.interpolate(disp_small_up, scale_factor=0.25, mode='bilinear', align_corners=True) * 0.25 # Init disp will be 1/4 + disp = self.forward(image1, image2, iters=iters, test_mode=test_mode, low_memory=low_memory, init_disp=init_disp) + disp = padder.unpad(disp) + return disp + +FoundationStereoLite = FastFoundationStereo + + +class TrtFeatureRunner(nn.Module): + def __init__(self, model): + super().__init__() + self.feature = model.feature + self.stem_2 = model.stem_2 + + def forward(self, image1, image2): + image1 = normalize_image(image1) + image2 = normalize_image(image2) + B = len(image1) + out = self.feature(torch.cat([image1, image2], dim=0)) + features_left = [o[:B] for o in out] + features_right = [o[B:] for o in out] + stem_2x = self.stem_2(image1) + return *features_left, features_right[0], stem_2x + + +class TrtPostRunner(nn.Module): + def __init__(self, model): + super().__init__() + self.args = model.args + self.dtype = model.dtype + self.register_buffer("dx", model.dx) + self.proj_cmb = model.proj_cmb + self.corr_stem = model.corr_stem + self.corr_feature_att = model.corr_feature_att + self.cost_agg = model.cost_agg + self.classifier = model.classifier + self.update_block = model.update_block + self.sam = model.sam + self.cam = model.cam + self.feature = model.feature + self.spx_2_gru = model.spx_2_gru + self.spx_gru = model.spx_gru + self.cnet = model.cnet + + def upsample_disp(self, disp, mask_feat_4, stem_2x): + with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE): + xspx = self.spx_2_gru(mask_feat_4, stem_2x) # 1/2 resolution + spx_pred = self.spx_gru(xspx) + spx_pred = F.softmax(spx_pred, 1) + up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1) + return up_disp.to(self.dtype) + + + def forward(self, features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x, gwc_volume): + features_left = [features_left_04, features_left_08, features_left_16, features_left_32] + left_tmp = self.proj_cmb(features_left_04) + right_tmp = self.proj_cmb(features_right_04) + concat_volume = build_concat_volume_optimized_pytorch(left_tmp, right_tmp, maxdisp=self.args.max_disp//4) + del left_tmp, right_tmp + comb_volume = torch.cat([gwc_volume, concat_volume], dim=1) + del concat_volume, gwc_volume + comb_volume = self.corr_stem(comb_volume) + comb_volume = self.corr_feature_att(comb_volume, features_left_04) + comb_volume = self.cost_agg(comb_volume, features_left) + + # Init disp from geometry encoding volume + logits = self.classifier(comb_volume).squeeze(1) + prob = F.softmax(logits, dim=1) + init_disp = disparity_regression(prob, self.args.max_disp//4) + + cnet_list = self.cnet(features_left[0], features_left[1], features_left[2]) + cnet_list = list(cnet_list) + net_list = [torch.tanh(x[0]) for x in cnet_list] + inp_list = [torch.relu(x[1]) for x in cnet_list] + inp_list = [self.cam(x) * x for x in inp_list] + att = [self.sam(x) for x in inp_list] + + geo_fn = Combined_Geo_Encoding_Volume(features_left_04.to(self.dtype), features_right_04.to(self.dtype), comb_volume.to(self.dtype), num_levels=self.args.corr_levels) + b, c, h, w = features_left[0].shape + coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1) + disp = init_disp.to(self.dtype) + + # GRUs iterations to update disparity (1/4 resolution) + for itr in range(self.args.valid_iters): + disp = disp.detach() + geo_feat = geo_fn(disp, coords, dx=self.dx, low_memory=True) + net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat.to(self.dtype), disp, att) + + disp = disp + delta_disp.to(self.dtype) + if itr < self.args.valid_iters-1: + continue + + disp_up = self.upsample_disp(disp.to(self.dtype), mask_feat_4.to(self.dtype), stem_2x.to(self.dtype)) + + return disp_up + + +class TrtRunner(nn.Module): + def __init__(self, args, feature_runner_engine_path, post_runner_engine_path): + super().__init__() + import tensorrt as trt + self.args = args + with open(feature_runner_engine_path, 'rb') as file: + engine_data = file.read() + self.trt_logger = trt.Logger(trt.Logger.WARNING) + self.feature_engine = trt.Runtime(self.trt_logger).deserialize_cuda_engine(engine_data) + self.feature_context = self.feature_engine.create_execution_context() + + with open(post_runner_engine_path, 'rb') as file: + engine_data = file.read() + self.post_engine = trt.Runtime(self.trt_logger).deserialize_cuda_engine(engine_data) + self.post_context = self.post_engine.create_execution_context() + self.max_disp = args.max_disp + self.cv_group = args.get('cv_group', 8) + + def trt_dtype_to_torch(self, dt): + import tensorrt as trt + if dt==trt.DataType.FLOAT: return torch.float32 + if dt==trt.DataType.HALF: return torch.float16 + if dt==trt.DataType.BF16: return torch.bfloat16 + if dt==trt.DataType.INT32: return torch.int32 + if dt==trt.DataType.INT8: return torch.int8 + if dt==trt.DataType.BOOL: return torch.bool + raise RuntimeError(f"Unsupported TRT dtype: {dt}") + + def get_io_tensor_names(self, engine, mode): + names = [] + n = engine.num_io_tensors + for i in range(n): + name = engine.get_tensor_name(i) + if engine.get_tensor_mode(name)==mode: + names.append(name) + return names + + def run_trt(self, engine, context, inputs_by_name:dict): + import tensorrt as trt + for name, tensor in list(inputs_by_name.items()): + expected_dtype = self.trt_dtype_to_torch(engine.get_tensor_dtype(name)) + if tensor.dtype != expected_dtype: inputs_by_name[name] = tensor.to(expected_dtype) + if not inputs_by_name[name].is_contiguous(): inputs_by_name[name] = inputs_by_name[name].contiguous() + context.set_input_shape(name, tuple(inputs_by_name[name].shape)) + outputs = {} + out_names = [n for n in self.get_io_tensor_names(engine, trt.TensorIOMode.OUTPUT)] + for name in out_names: + shp = tuple(context.get_tensor_shape(name)) + dtype = self.trt_dtype_to_torch(engine.get_tensor_dtype(name)) + outputs[name] = torch.empty(shp, device='cuda', dtype=dtype) + for name, tensor in inputs_by_name.items(): context.set_tensor_address(name, int(tensor.data_ptr())) + for name, tensor in outputs.items(): context.set_tensor_address(name, int(tensor.data_ptr())) + stream = torch.cuda.current_stream().cuda_stream + ok = context.execute_async_v3(stream) + assert ok + return outputs + + def forward(self, image1, image2): + import tensorrt as trt + feat_out = self.run_trt(self.feature_engine, self.feature_context, {'left': image1, 'right': image2}) + gwc_volume = build_gwc_volume_triton(feat_out['features_left_04'].half(), feat_out['features_right_04'].half(), self.args.max_disp//4, self.cv_group, normalize=self.args.normalize) + post_inputs = feat_out + post_inputs['gwc_volume'] = gwc_volume + in_names = self.get_io_tensor_names(self.post_engine, trt.TensorIOMode.INPUT) + tmp_keys = list(post_inputs.keys()) + for k in tmp_keys: + if k not in in_names: + del post_inputs[k] + out = self.run_trt(self.post_engine, self.post_context, post_inputs) + disp = out['disp'] + return disp \ No newline at end of file diff --git a/cv/fast-foundationstereo/core/geometry.py b/cv/fast-foundationstereo/core/geometry.py new file mode 100755 index 0000000..80762c1 --- /dev/null +++ b/cv/fast-foundationstereo/core/geometry.py @@ -0,0 +1,81 @@ +import torch,os,sys +import torch.nn.functional as F +from core.utils.utils import bilinear_sampler, bilinear_sampler1d +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') + +class Combined_Geo_Encoding_Volume: + def __init__(self, init_fmap1, init_fmap2, geo_volume, num_levels=2): + self.num_levels = num_levels + self.geo_volume_pyramid = [] + self.init_corr_pyramid = [] + + # all pairs correlation + init_corr = Combined_Geo_Encoding_Volume.corr(init_fmap1, init_fmap2) + + b, h, w, _, w2 = init_corr.shape + b, c, d, h, w = geo_volume.shape + geo_volume = geo_volume.permute(0, 3, 4, 1, 2).reshape(b*h*w, c, 1, d) + + init_corr = init_corr.view(b*h*w, 1, 1, w2) + self.geo_volume_pyramid.append(geo_volume) + self.init_corr_pyramid.append(init_corr) + for _ in range(self.num_levels-1): + geo_volume = F.avg_pool2d(geo_volume, [1,2], stride=[1,2]) + self.geo_volume_pyramid.append(geo_volume) + + for _ in range(self.num_levels-1): + init_corr = F.avg_pool2d(init_corr, [1,2], stride=[1,2]) + self.init_corr_pyramid.append(init_corr) + + + + def __call__(self, disp, coords, dx, low_memory=True): + b, _, h, w = disp.shape + out_pyramid = [] + for i in range(self.num_levels): + with torch.profiler.record_function(f"make disp_lvl {i}"): + geo_volume = self.geo_volume_pyramid[i] + x0 = dx + disp.view(b*h*w, 1, 1, 1) / 2**i + with torch.profiler.record_function(f"bilinear_sampler geo_volume {i}"): + if low_memory: + geo_volume = bilinear_sampler1d(geo_volume, x0, mode='bilinear', align_corners=True) + else: + y0 = torch.zeros_like(x0) + disp_lvl = torch.cat([x0,y0], dim=-1) + geo_volume = bilinear_sampler(geo_volume, disp_lvl, low_memory=low_memory) + geo_volume = geo_volume.view(b, h, w, -1) #(b, h, h, 3x3xC) + + with torch.profiler.record_function(f"make init_coords_lvl {i}"): + init_corr = self.init_corr_pyramid[i] # (B*H*W, 1, 1, W2) + init_x0 = coords.view(b*h*w, 1, 1, 1)/2**i - disp.view(b*h*w, 1, 1, 1) / 2**i + dx # X on right image + with torch.profiler.record_function(f"bilinear_sampler init_corr {i}"): + if low_memory: + init_corr = bilinear_sampler1d(init_corr, init_x0, mode='bilinear', align_corners=True) + else: + init_coords_lvl = torch.cat([init_x0,y0], dim=-1) + init_corr = bilinear_sampler(init_corr, init_coords_lvl, low_memory=low_memory) + init_corr = init_corr.view(b, h, w, -1) + + out_pyramid.append(geo_volume) + out_pyramid.append(init_corr) + + with torch.profiler.record_function(f"make out_pyramid"): + out_pyramid = torch.cat(out_pyramid, dim=-1) + return out_pyramid.permute(0, 3, 1, 2) #(B,C,H,W) + + + @staticmethod + def corr(fmap1, fmap2, normalize=True): + with torch.profiler.record_function("build corr"): + B, D, H, W1 = fmap1.shape + _, _, _, W2 = fmap2.shape + fmap1 = fmap1.view(B, D, H, W1) + fmap2 = fmap2.view(B, D, H, W2) + if normalize: + with torch.cuda.amp.autocast(enabled=False): + corr = torch.einsum('aijk,aijh->ajkh', F.normalize(fmap1.float(), dim=1), F.normalize(fmap2.float(), dim=1)) + else: + corr = corr.view(B, H, W1, 1, W2).to(fmap1.dtype) + corr = corr.view(B, H, W1, 1, W2).to(fmap1.dtype) + return corr \ No newline at end of file diff --git a/cv/fast-foundationstereo/core/submodule.py b/cv/fast-foundationstereo/core/submodule.py new file mode 100755 index 0000000..6764d64 --- /dev/null +++ b/cv/fast-foundationstereo/core/submodule.py @@ -0,0 +1,675 @@ +import torch,pdb,os,sys +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from Utils import AMP_DTYPE +import Utils as U +try: + import triton + import triton.language as tl +except Exception: + triton = None + tl = None + +def _is_contiguous(tensor: torch.Tensor) -> bool: + if torch.jit.is_scripting(): + return tensor.is_contiguous() + else: + return tensor.is_contiguous(memory_format=torch.contiguous_format) + + + +class LayerNorm2d(nn.LayerNorm): + r""" https://huggingface.co/spaces/Roll20/pet_score/blob/b258ef28152ab0d5b377d9142a23346f863c1526/lib/timm/models/convnext.py#L85 + LayerNorm for channels_first tensors with 2d spatial dimensions (ie N, C, H, W). + """ + + def __init__(self, normalized_shape, eps=1e-6): + """ + @normalized_shape: channel dim + """ + super().__init__(normalized_shape, eps=eps) + + def forward(self, x) -> torch.Tensor: + """ + @x: (B,C,H,W) + """ + if _is_contiguous(x): + return F.layer_norm(x.permute(0, 2, 3, 1), self.normalized_shape, self.weight, self.bias, self.eps).permute(0, 3, 1, 2).contiguous() + else: + s, u = torch.var_mean(x, dim=1, keepdim=True) + x = (x - u) * torch.rsqrt(s + self.eps) + x = x * self.weight[:, None, None] + self.bias[:, None, None] + return x + + +class BasicConv(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, norm='batch', **kwargs): + super(BasicConv, self).__init__() + + self.relu = nn.LeakyReLU(inplace=True) if relu else nn.Identity() + self.use_bn = bn + self.bn = nn.Identity() + if is_3d: + if deconv: + self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) + if self.use_bn: + if norm=='batch': + self.bn = nn.BatchNorm3d(out_channels) + elif norm=='instance': + self.bn = nn.InstanceNorm3d(out_channels) + else: + if deconv: + self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + if self.use_bn: + if norm=='batch': + self.bn = nn.BatchNorm2d(out_channels) + elif norm=='instance': + self.bn = nn.InstanceNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + if self.use_bn: + x = self.bn(x) + if isinstance(self.relu, bool): + if self.relu: + self.relu = nn.LeakyReLU(inplace=True) + else: + self.relu = nn.Identity() + x = self.relu(x) + return x + + +class Conv3dNormActReduced(nn.Module): + def __init__(self, C_in, C_out, hidden=None, kernel_size=3, kernel_disp=None, stride=1, norm=nn.BatchNorm3d): + super().__init__() + if kernel_disp is None: + kernel_disp = kernel_size + if hidden is None: + hidden = C_out + self.conv1 = nn.Sequential( + nn.Conv3d(C_in, hidden, kernel_size=(1,kernel_size,kernel_size), padding=(0, kernel_size//2, kernel_size//2), stride=(1, stride, stride)), + norm(hidden), + nn.ReLU(), + ) + self.conv2 = nn.Sequential( + nn.Conv3d(hidden, C_out, kernel_size=(kernel_disp, 1, 1), padding=(kernel_disp//2, 0, 0), stride=(stride, 1, 1)), + norm(C_out), + nn.ReLU(), + ) + + + def forward(self, x): + """ + @x: (B,C,D,H,W) + """ + x = self.conv1(x) + x = self.conv2(x) + return x + + +class ResnetBasicBlock(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm2d, bias=False): + super().__init__() + self.norm_layer = norm_layer + if groups != 1 or base_width != 64: + raise ValueError('BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding) + if self.norm_layer is not None: + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, bias=bias, padding=padding) + if self.norm_layer is not None: + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + + def forward(self, x): + identity = x + + out = self.conv1(x) + if self.norm_layer is not None: + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + if self.norm_layer is not None: + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + out += identity + out = self.relu(out) + + return out + + +class ResnetBasicBlock3D(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm3d, bias=False): + super().__init__() + self.norm_layer = norm_layer + if groups != 1 or base_width != 64: + raise ValueError('BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding) + if self.norm_layer is not None: + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv3d(planes, planes, kernel_size=kernel_size, bias=bias, padding=padding) + if self.norm_layer is not None: + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + + def forward(self, x): + identity = x + + out = self.conv1(x) + if self.norm_layer is not None: + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + if self.norm_layer is not None: + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + out += identity + out = self.relu(out) + + return out + + +class FlashMultiheadAttention(nn.Module): + def __init__(self, embed_dim, num_heads): + super().__init__() + self.num_heads = num_heads + self.embed_dim = embed_dim + self.head_dim = embed_dim // num_heads + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + + self.q_proj = nn.Linear(embed_dim, embed_dim) + self.k_proj = nn.Linear(embed_dim, embed_dim) + self.v_proj = nn.Linear(embed_dim, embed_dim) + self.out_proj = nn.Linear(embed_dim, embed_dim) + + def forward(self, query, key, value, attn_mask=None, window_size=(-1,-1)): + """ + @query: (B,L,C) + """ + B,L,C = query.shape + Q = self.q_proj(query) + K = self.k_proj(key) + V = self.v_proj(value) + + Q = Q.view(Q.size(0), Q.size(1), self.num_heads, self.head_dim) + K = K.view(K.size(0), K.size(1), self.num_heads, self.head_dim) + V = V.view(V.size(0), V.size(1), self.num_heads, self.head_dim) + + attn_output = F.scaled_dot_product_attention(Q, K, V) + + attn_output = attn_output.reshape(B,L,-1) + output = self.out_proj(attn_output) + + return output + + + +class FlashAttentionTransformerEncoderLayer(nn.Module): + def __init__(self, embed_dim, num_heads, dim_feedforward, dropout=0.1, act=nn.GELU, norm=nn.LayerNorm): + super().__init__() + self.self_attn = FlashMultiheadAttention(embed_dim, num_heads) + self.act = act() + + self.linear1 = nn.Linear(embed_dim, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, embed_dim) + + self.norm1 = norm(embed_dim) + self.norm2 = norm(embed_dim) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + def forward(self, src, src_mask=None, window_size=(-1, -1)): + dtype = src.dtype + src2 = self.self_attn(src, src, src, src_mask, window_size=window_size) + src = src + self.dropout1(src2) + src = self.norm1(src).to(dtype) + + src2 = self.linear2(self.dropout(self.act(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src).to(dtype) + + return src + + +class Conv2x(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, relu=True, keep_dispc=False): + super(Conv2x, self).__init__() + self.concat = concat + self.is_3d = is_3d + if deconv and is_3d: + kernel = (4, 4, 4) + elif deconv: + kernel = 4 + else: + kernel = 3 + + if deconv and is_3d and keep_dispc: + kernel = (1, 4, 4) + stride = (1, 2, 2) + padding = (0, 1, 1) + self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=bn, relu=True, kernel_size=kernel, stride=stride, padding=padding) + else: + self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=bn, relu=True, kernel_size=kernel, stride=2, padding=1) + + if self.concat: + mul = 2 if keep_concat else 1 + self.conv2 = BasicConv(out_channels*2, out_channels*mul, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1) + else: + self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1) + + def forward(self, x, rem): + x = self.conv1(x) + if x.shape != rem.shape: + x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='bilinear') + if self.concat: + x = torch.cat((x, rem), 1) + else: + x = x + rem + x = self.conv2(x) + return x + + +class BasicConv_IN(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, IN=True, relu=True, **kwargs): + super(BasicConv_IN, self).__init__() + if relu: + self.relu = nn.LeakyReLU(inplace=True) + else: + self.relu = nn.Identity() + self.use_in = IN + if is_3d: + if deconv: + self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) + self.IN = nn.InstanceNorm3d(out_channels) + else: + if deconv: + self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.IN = nn.InstanceNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + if self.use_in: + x = self.IN(x) + if isinstance(self.relu, bool): + if self.relu: + self.relu = nn.LeakyReLU(inplace=True) + else: + self.relu = nn.Identity() + x = self.relu(x) + return x + + +class Conv2x_IN(nn.Module): + def __init__(self, in_channels, out_channels, c_middle=None, deconv=False, is_3d=False, concat=True, keep_concat=True, IN=True, relu=True, keep_dispc=False): + super(Conv2x_IN, self).__init__() + self.concat = concat + self.is_3d = is_3d + if deconv and is_3d: + kernel = (4, 4, 4) + elif deconv: + kernel = 4 + else: + kernel = 3 + if c_middle is None: + c_middle = out_channels + + if deconv and is_3d and keep_dispc: + kernel = (1, 4, 4) + stride = (1, 2, 2) + padding = (0, 1, 1) + self.conv1 = BasicConv_IN(in_channels, c_middle, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=stride, padding=padding) + else: + self.conv1 = BasicConv_IN(in_channels, c_middle, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=2, padding=1) + + if self.concat: + mul = 2 if keep_concat else 1 + self.conv2 = ResnetBasicBlock(out_channels*2, out_channels*mul, kernel_size=3, stride=1, padding=1, norm_layer=nn.InstanceNorm2d) + else: + self.conv2 = BasicConv_IN(c_middle, out_channels, False, is_3d, IN, relu, kernel_size=3, stride=1, padding=1) + + def forward(self, x, rem): + x = self.conv1(x) + if x.shape != rem.shape: + x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='bilinear') + if self.concat: + x = torch.cat((x, rem), 1) + else: + x = x + rem + x = self.conv2(x) + return x + + + +@torch.compile +def build_gwc_volume_optimized_pytorch1(refimg_fea: torch.Tensor, targetimg_fea: torch.Tensor, maxdisp: int, num_groups: int, normalize=True): + dtype = refimg_fea.dtype + B, C, H, W = refimg_fea.shape + channels_per_group = C // num_groups + + ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W) + padded_target = F.pad(targetimg_fea, (maxdisp - 1, 0, 0, 0)) + unfolded_target = padded_target.unfold(3, W, 1) + target_volume = torch.flip(unfolded_target, [3]).permute(0, 1, 3, 2, 4) + ref_volume = ref_volume.view(B, num_groups, channels_per_group, maxdisp, H, W) + target_volume = target_volume.view(B, num_groups, channels_per_group, maxdisp, H, W) + if normalize: + ref_volume = F.normalize(ref_volume.float(), dim=2).to(dtype) + target_volume = F.normalize(target_volume.float(), dim=2).to(dtype) + + cost_volume = (ref_volume * target_volume).sum(dim=2) + + return cost_volume.contiguous() + + +if triton is not None and torch.cuda.is_available(): + @triton.autotune(configs=[ + triton.Config({'BLOCK_C':4,'BLOCK_W':128,'BLOCK_D':8}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_C':8,'BLOCK_W':128,'BLOCK_D':8}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_C':16,'BLOCK_W':128,'BLOCK_D':8}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_C':64,'BLOCK_W':128,'BLOCK_D':8}, num_warps=8, num_stages=2), + triton.Config({'BLOCK_C':128,'BLOCK_W':64,'BLOCK_D':8}, num_warps=8, num_stages=2), + triton.Config({'BLOCK_C':128,'BLOCK_W':128,'BLOCK_D':8}, num_warps=8, num_stages=2), + ], key=['C','W','D','G','K','NORMALIZE']) + @triton.jit + def _gwc_triton_kernel(ref_ptr, tar_ptr, ref_norm_ptr, tar_norm_ptr, out_ptr, BH, C, W, D: tl.constexpr, G: tl.constexpr, K: tl.constexpr, + stride_rn, stride_rw, stride_rc, stride_tn, stride_tw, stride_tc, + stride_nn, stride_ng, stride_nw, + stride_on, stride_og, stride_od, stride_ow, + NORMALIZE: tl.constexpr, + BLOCK_C: tl.constexpr, BLOCK_W: tl.constexpr, BLOCK_D: tl.constexpr): + pid0 = tl.program_id(0) + db = tl.program_id(1) + wb = tl.program_id(2) + bh = pid0 // G + g = pid0 % G + w_off = wb*BLOCK_W + tl.arange(0, BLOCK_W) + d_off = db*BLOCK_D + tl.arange(0, BLOCK_D) + w_mask = w_off < W + w_src = w_off[None, :] - d_off[:, None] + td_mask = (w_src >= 0) & w_mask[None, :] + acc = tl.zeros((BLOCK_D, BLOCK_W), dtype=tl.float32) + for k0 in tl.static_range(0, K, BLOCK_C): + k_off = k0 + tl.arange(0, BLOCK_C) + k_mask = k_off < K + c_idx = g*K + k_off + ref_ptrs = ref_ptr + bh*stride_rn + w_off[None, :]*stride_rw + c_idx[:, None]*stride_rc + ref_vals = tl.load(ref_ptrs, mask=k_mask[:, None] & w_mask[None, :], other=0.).to(tl.float32) + tar_ptrs = tar_ptr + bh*stride_tn + w_src[None, :, :]*stride_tw + c_idx[:, None, None]*stride_tc + tar_vals = tl.load(tar_ptrs, mask=k_mask[:, None, None] & td_mask[None, :, :], other=0.).to(tl.float32) + acc += tl.sum(tar_vals * ref_vals[:, None, :], axis=0) + + if NORMALIZE: + norm_offset = bh*stride_nn + g*stride_ng + ref_norm = tl.load(ref_norm_ptr + norm_offset + w_off*stride_nw, mask=w_mask, other=1.0).to(tl.float32) + tar_norm = tl.load(tar_norm_ptr + norm_offset + w_src*stride_nw, mask=td_mask, other=1.0).to(tl.float32) + denom = (ref_norm[None, :] * tar_norm) + 1e-5 + acc = acc / denom + out_ptrs = out_ptr + bh*stride_on + g*stride_og + d_off[:, None]*stride_od + w_off[None, :]*stride_ow + tl.store(out_ptrs, acc, mask=w_mask[None, :]) + +@torch.no_grad() +def build_gwc_volume_triton(refimg_fea: torch.Tensor, targetimg_fea: torch.Tensor, maxdisp: int, num_groups: int, normalize=True): + if triton is None: + raise RuntimeError('Triton is not available. Please install triton to use build_gwc_volume_triton.') + B, C, H, W = refimg_fea.shape + assert maxdisp > 0 and C % num_groups == 0 + K = C // num_groups + in_dtype = refimg_fea.dtype if refimg_fea.dtype in (torch.float16, torch.bfloat16, torch.float32) else torch.float32 + + if normalize: + ref_norm = refimg_fea.float().view(B, num_groups, K, H, W).norm(dim=2) + tar_norm = targetimg_fea.float().view(B, num_groups, K, H, W).norm(dim=2) + ref_norm = ref_norm.permute(0, 2, 1, 3).reshape(B*H, num_groups, W).to(in_dtype).contiguous() + tar_norm = tar_norm.permute(0, 2, 1, 3).reshape(B*H, num_groups, W).to(in_dtype).contiguous() + else: + # Dummy tensors; kernel won't read them when NORMALIZE=False + ref_norm = refimg_fea.new_empty((1, 1, 1), dtype=in_dtype) + tar_norm = refimg_fea.new_empty((1, 1, 1), dtype=in_dtype) + + ref = refimg_fea.to(in_dtype) + tar = targetimg_fea.to(in_dtype) + ref_bhwc = ref.permute(0, 2, 3, 1).view(B * H, W, C).contiguous() + tar_bhwc = tar.permute(0, 2, 3, 1).view(B * H, W, C).contiguous() + out_bhw = torch.empty((B * H, num_groups, maxdisp, W), device=ref.device, dtype=in_dtype) + BH = B * H + D_eff = min(maxdisp, W) + grid = lambda META: (BH * num_groups, triton.cdiv(D_eff, META['BLOCK_D']), triton.cdiv(W, META['BLOCK_W'])) + _gwc_triton_kernel[grid](ref_bhwc, tar_bhwc, ref_norm, tar_norm, out_bhw, BH, C, W, D_eff, num_groups, K, + ref_bhwc.stride(0), ref_bhwc.stride(1), ref_bhwc.stride(2), + tar_bhwc.stride(0), tar_bhwc.stride(1), tar_bhwc.stride(2), + ref_norm.stride(0), ref_norm.stride(1), ref_norm.stride(2), + out_bhw.stride(0), out_bhw.stride(1), out_bhw.stride(2), out_bhw.stride(3), + NORMALIZE=normalize) + if D_eff < maxdisp: out_bhw[:, :, D_eff:, :] = 0 + volume = out_bhw.view(B, H, num_groups, maxdisp, W).permute(0, 2, 3, 1, 4).contiguous() + return volume + + + +@torch.compile +def build_concat_volume_optimized_pytorch(refimg_fea, targetimg_fea, maxdisp:int): + B, C, H, W = refimg_fea.shape + ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W) + shifted_target_list = [F.pad(targetimg_fea, (int(d), 0, 0, 0), "constant", 0.0)[:, :, :, :W] for d in range(maxdisp)] + target_volume = torch.stack(shifted_target_list, dim=2) + volume = torch.cat((ref_volume, target_volume), dim=1) + return volume.contiguous() + + +@torch.compile +def build_concat_volume_optimized_pytorch1(refimg_fea, targetimg_fea, maxdisp:int): + B, C, H, W = refimg_fea.shape + + ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W) + padded_target = F.pad(targetimg_fea, (maxdisp - 1, 0, 0, 0)) # (B, C, H, W + maxdisp - 1) + unfolded_target = padded_target.unfold(dimension=3, size=W, step=1) # (B, C, H, maxdisp, W) + target_volume = torch.flip(unfolded_target, [3]).permute(0, 1, 3, 2, 4) + volume = torch.cat((ref_volume, target_volume), dim=1) + return volume.contiguous() + + + + +def disparity_regression(x, maxdisp): + assert len(x.shape) == 4 + disp_values = torch.arange(0, maxdisp, dtype=x.dtype, device=x.device) + disp_values = disp_values.reshape(1, maxdisp, 1, 1) + return torch.sum(x * disp_values, 1, keepdim=True) #(B,1,H,W) + + +class FeatureAtt(nn.Module): + def __init__(self, cv_chan, feat_chan): + super(FeatureAtt, self).__init__() + + self.feat_att = nn.Sequential( + BasicConv(feat_chan, feat_chan//2, kernel_size=1, stride=1, padding=0), + nn.Conv2d(feat_chan//2, cv_chan, 1) + ) + + def forward(self, cv, feat): + ''' + @cv: cost volume (B,C,D,H,W) + @feat: (B,C,H,W) + ''' + feat_att = self.feat_att(feat).unsqueeze(2) #(B,C,1,H,W) + cv = torch.sigmoid(feat_att)*cv + return cv + +def context_upsample(disp_low, up_weights): + """ + @disp_low: (b,1,h,w) 1/4 resolution + @up_weights: (b,9,4*h,4*w) Image resolution + """ + b, c, h, w = disp_low.shape + + disp_unfold = F.unfold(disp_low.reshape(b,c,h,w),3,1,1).reshape(b,-1,h,w) + disp_unfold = F.interpolate(disp_unfold,(h*4,w*4),mode='nearest').reshape(b,9,h*4,w*4) + + disp = (disp_unfold*up_weights).sum(1) + + return disp + + + +class PositionalEmbedding(nn.Module): + def __init__(self, d_model, max_len=512): + super().__init__() + + # Compute the positional encodings once in log space. + pe = torch.zeros(max_len, d_model, dtype=torch.float) + pe.require_grad = False + + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) #(N,1) + div_term = (torch.arange(0, d_model, 2, dtype=torch.float) * -(np.log(10000.0) / d_model)).exp()[None] + + pe[:, 0::2] = torch.sin(position * div_term) #(N, d_model/2) + pe[:, 1::2] = torch.cos(position * div_term) + + pe = pe.unsqueeze(0) + self.pe = pe + + + def forward(self, x, resize_embed=False): + ''' + @x: (B,N,D) + ''' + dtype = x.dtype + self.pe = self.pe.to(x.device).to(x.dtype) + pe = self.pe + if pe.shape[1] 0 else None + + def forward(self, x): + input = x + x = self.dwconv(x) + x = self.norm(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + x + return x + diff --git a/cv/fast-foundationstereo/core/update.py b/cv/fast-foundationstereo/core/update.py new file mode 100755 index 0000000..ed641aa --- /dev/null +++ b/cv/fast-foundationstereo/core/update.py @@ -0,0 +1,108 @@ +import torch,os,sys +import torch.nn as nn +import torch.nn.functional as F +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from core.submodule import EdgeNextConvEncoder + + +class DispHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256, output_dim=1): + super(DispHead, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(input_dim, input_dim, kernel_size=3, padding=1), + nn.ReLU(), + EdgeNextConvEncoder(input_dim, expan_ratio=4, kernel_size=7, norm=None), + EdgeNextConvEncoder(input_dim, expan_ratio=4, kernel_size=7, norm=None), + nn.Conv2d(input_dim, output_dim, 3, padding=1), + ) + + def forward(self, x): + return self.conv(x) + + +class BasicMotionEncoder(nn.Module): + def __init__(self, args, ngroup=8): + super(BasicMotionEncoder, self).__init__() + self.args = args + cor_planes = args.corr_levels * (2*args.corr_radius + 1) * (ngroup+1) + self.convc1 = nn.Conv2d(cor_planes, 256, kernel_size=1, padding=0) + self.convc2 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.convd1 = nn.Conv2d(1, 64, kernel_size=7, padding=3) + self.convd2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) + self.conv = nn.Conv2d(64+256, args.hidden_dims[0]-1, kernel_size=1, padding=0) + + def forward(self, disp, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + disp_ = F.relu(self.convd1(disp)) + disp_ = F.relu(self.convd2(disp_)) + + cor_disp = torch.cat([cor, disp_], dim=1) + out = F.relu(self.conv(cor_disp)) + return torch.cat([out, disp], dim=1) + +class RaftConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=256, kernel_size=3): + super().__init__() + self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2) + self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2) + self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2) + + def forward(self, h, x, hx): + z = torch.sigmoid(self.convz(hx)) + r = torch.sigmoid(self.convr(hx)) + q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + return h + + +class SelectiveConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=256, small_kernel_size=1, large_kernel_size=3, patch_size=None): + super(SelectiveConvGRU, self).__init__() + self.conv0 = nn.Sequential( + nn.Conv2d(input_dim, input_dim, kernel_size=3, padding=1), + nn.ReLU(), + ) + self.conv1 = nn.Sequential( + nn.Conv2d(input_dim+hidden_dim, input_dim+hidden_dim, kernel_size=3, padding=1), + nn.ReLU(), + ) + self.small_gru = RaftConvGRU(hidden_dim, input_dim, small_kernel_size) + self.large_gru = RaftConvGRU(hidden_dim, input_dim, large_kernel_size) + + def forward(self, att, h, *x): + x = torch.cat(x, dim=1) + x = self.conv0(x) + hx = torch.cat([x, h], dim=1) + hx = self.conv1(hx) + h = self.small_gru(h, x, hx) * att + self.large_gru(h, x, hx) * (1 - att) + + return h + + +class BasicSelectiveMultiUpdateBlock(nn.Module): + def __init__(self, args, hidden_dim=128, volume_dim=8): + super().__init__() + self.args = args + self.encoder = BasicMotionEncoder(args, volume_dim) + + self.gru04 = SelectiveConvGRU(hidden_dim, hidden_dim*2) + self.disp_head = DispHead(hidden_dim, 256) + self.mask = nn.Sequential( + nn.Conv2d(hidden_dim, 64, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(64, 32, 3, padding=1), + nn.ReLU(inplace=True), + ) + + def forward(self, net, inp, corr, disp, att): + motion_features = self.encoder(disp, corr) + motion_features = torch.cat([inp[0], motion_features], dim=1) + net[0] = self.gru04(att[0], net[0], motion_features) + + delta_disp = self.disp_head(net[0]) + + mask = .25 * self.mask(net[0]) + return net, mask, delta_disp + diff --git a/cv/fast-foundationstereo/core/utils/__init__.py b/cv/fast-foundationstereo/core/utils/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/cv/fast-foundationstereo/core/utils/frame_utils.py b/cv/fast-foundationstereo/core/utils/frame_utils.py new file mode 100755 index 0000000..6f0695d --- /dev/null +++ b/cv/fast-foundationstereo/core/utils/frame_utils.py @@ -0,0 +1,202 @@ +import numpy as np +from PIL import Image +from os.path import basename, exists, splitext +import os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../../') +import re +import json +import imageio +import cv2 +from turbojpeg import TurboJPEG, TJPF_GRAY, TJSAMP_GRAY, TJFLAG_PROGRESSIVE, TJFLAG_FASTUPSAMPLE, TJFLAG_FASTDCT +jpeg = TurboJPEG() +cv2.setNumThreads(0) +cv2.ocl.setUseOpenCL(False) + +TAG_CHAR = np.array([202021.25], np.float32) + +def readFlow(fn): + """ Read .flo file in Middlebury format""" + # Code adapted from: + # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy + + # WARNING: this will work on little-endian architectures (eg Intel x86) only! + # print 'fn = %s'%(fn) + with open(fn, 'rb') as f: + magic = np.fromfile(f, np.float32, count=1) + if 202021.25 != magic: + print('Magic number incorrect. Invalid .flo file') + return None + else: + w = np.fromfile(f, np.int32, count=1) + h = np.fromfile(f, np.int32, count=1) + # print 'Reading %d x %d flo file\n' % (w, h) + data = np.fromfile(f, np.float32, count=2*int(w)*int(h)) + # Reshape data into 3D array (columns, rows, bands) + # The reshape here is for visualization, the original code is (w,h,2) + return np.resize(data, (int(h), int(w), 2)) + +def readPFM(file): + file = open(file, 'rb') + + color = None + width = None + height = None + scale = None + endian = None + + header = file.readline().rstrip() + if header == b'PF': + color = True + elif header == b'Pf': + color = False + else: + raise Exception('Not a PFM file.') + + dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline()) + if dim_match: + width, height = map(int, dim_match.groups()) + else: + raise Exception('Malformed PFM header.') + + scale = float(file.readline().rstrip()) + if scale < 0: # little-endian + endian = '<' + scale = -scale + else: + endian = '>' # big-endian + + data = np.fromfile(file, endian + 'f') + shape = (height, width, 3) if color else (height, width) + + data = np.reshape(data, shape) + data = np.flipud(data) + return data + +def writePFM(file, array): + import os + assert type(file) is str and type(array) is np.ndarray and \ + os.path.splitext(file)[1] == ".pfm" + with open(file, 'wb') as f: + H, W = array.shape + headers = ["Pf\n", f"{W} {H}\n", "-1\n"] + for header in headers: + f.write(str.encode(header)) + array = np.flip(array, axis=0).astype(np.float32) + f.write(array.tobytes()) + + + +def writeFlow(filename,uv,v=None): + """ Write optical flow to file. + + If v is None, uv is assumed to contain both u and v channels, + stacked in depth. + Original code by Deqing Sun, adapted from Daniel Scharstein. + """ + nBands = 2 + + if v is None: + assert(uv.ndim == 3) + assert(uv.shape[2] == 2) + u = uv[:,:,0] + v = uv[:,:,1] + else: + u = uv + + assert(u.shape == v.shape) + height,width = u.shape + f = open(filename,'wb') + # write the header + f.write(TAG_CHAR) + np.array(width).astype(np.int32).tofile(f) + np.array(height).astype(np.int32).tofile(f) + # arrange into matrix form + tmp = np.zeros((height, width*nBands)) + tmp[:,np.arange(width)*2] = u + tmp[:,np.arange(width)*2 + 1] = v + tmp.astype(np.float32).tofile(f) + f.close() + + +def readFlowKITTI(filename): + flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR) + flow = flow[:,:,::-1].astype(np.float32) + flow, valid = flow[:, :, :2], flow[:, :, 2] + flow = (flow - 2**15) / 64.0 + return flow, valid + +def readDispKITTI(filename): + disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0 + valid = disp > 0.0 + return disp, valid + +# Method taken from /n/fs/raft-depth/RAFT-Stereo/datasets/SintelStereo/sdk/python/sintel_io.py +def readDispSintelStereo(file_name): + a = np.array(Image.open(file_name)) + d_r, d_g, d_b = np.split(a, axis=2, indices_or_sections=3) + disp = (d_r * 4 + d_g / (2**6) + d_b / (2**14))[..., 0] + mask = np.array(Image.open(file_name.replace('disparities', 'occlusions'))) + valid = ((mask == 0) & (disp > 0)) + return disp, valid + +# Method taken from https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt +def readDispFallingThings(file_name): + a = np.array(Image.open(file_name)) + with open('/'.join(file_name.split('/')[:-1] + ['_camera_settings.json']), 'r') as f: + intrinsics = json.load(f) + fx = intrinsics['camera_settings'][0]['intrinsic_settings']['fx'] + disp = (fx * 6.0 * 100) / a.astype(np.float32) + valid = disp > 0 + return disp, valid + +# Method taken from https://github.com/castacks/tartanair_tools/blob/master/data_type.md +def readDispTartanAir(file_name): + depth = np.load(file_name) + disp = 80.0 / depth + valid = disp > 0 + return disp, valid + + +def readDispMiddlebury(file_name): + assert basename(file_name) == 'disp0GT.pfm' + disp = readPFM(file_name).astype(np.float32) + assert len(disp.shape) == 2 + nocc_pix = file_name.replace('disp0GT.pfm', 'mask0nocc.png') + assert exists(nocc_pix) + nocc_pix = imageio.imread(nocc_pix) == 255 + assert np.any(nocc_pix) + return disp, nocc_pix + +def writeFlowKITTI(filename, uv): + uv = 64.0 * uv + 2**15 + valid = np.ones([uv.shape[0], uv.shape[1], 1]) + uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16) + cv2.imwrite(filename, uv[..., ::-1]) + + +def read_gen(file_name, pil=False): + ext = splitext(file_name)[-1] + if ext in ['.jpeg','.jpg']: + with open(file_name, 'rb') as ff: + bgr_array = jpeg.decode(ff.read()) + img = bgr_array[...,::-1] + return img + elif ext == '.png' or ext == '.ppm': + img = cv2.imread(file_name)[...,:3] + if len(img.shape)==3: + img = img[...,::-1] + elif len(img.shape)==2: + img = np.tile(img[...,None], (1,1,3)) + return img + elif ext == '.bin' or ext == '.raw': + return np.load(file_name) + elif ext == '.flo': + return readFlow(file_name).astype(np.float32) + elif ext == '.pfm': + flow = readPFM(file_name).astype(np.float32) + if len(flow.shape) == 2: + return flow + else: + return flow[:, :, :-1] + return [] \ No newline at end of file diff --git a/cv/fast-foundationstereo/core/utils/utils.py b/cv/fast-foundationstereo/core/utils/utils.py new file mode 100755 index 0000000..add2c8f --- /dev/null +++ b/cv/fast-foundationstereo/core/utils/utils.py @@ -0,0 +1,122 @@ +import torch,os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../../') +import torch.nn.functional as F +import numpy as np + + +class InputPadder: + """ Pads images such that dimensions are divisible by 8 """ + def __init__(self, dims, mode='sintel', divis_by=8, force_square=False): + self.ht, self.wd = dims[-2:] + if force_square: + max_side = max(self.ht, self.wd) + pad_ht = ((max_side // divis_by) + 1) * divis_by - self.ht + pad_wd = ((max_side // divis_by) + 1) * divis_by - self.wd + else: + pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by + pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by + if mode == 'sintel': + self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2] + else: + self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht] + + def pad(self, *inputs): + assert all((x.ndim == 4) for x in inputs) + return [F.pad(x, self._pad, mode='replicate') for x in inputs] + + def unpad(self, x): + assert x.ndim == 4 + ht, wd = x.shape[-2:] + c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]] + return x[..., c[0]:c[1], c[2]:c[3]] + + +@torch.compile +def bilinear_sampler1d(img, x_coords, mode='bilinear', align_corners=True): + """ + 1D bilinear sampling along width dimension only (for stereo applications) + Much faster than grid_sample for stereo where y is constant + + Args: + img: (B, C, 1, W) input tensor + x_coords: (B, 1, W_out, 1) x coordinates in pixel space [0, W-1] + mode: interpolation mode ('bilinear' or 'nearest') + align_corners: if True, corner pixels are aligned (like grid_sample) + + Returns: + sampled: (B, C, 1, W_coords) sampled tensor + mask: (B, 1, H, W) validity mask (if mask=True) + """ + B, C, H_img, W = img.shape + x = x_coords.reshape(B,-1) # (B, W_out) + + if align_corners: + # align_corners=True: coordinate range [0, W-1] maps to pixel centers + # This matches grid_sample with align_corners=True behavior + x_normalized = x + else: + # align_corners=False: coordinate range [0, W-1] maps to pixel edges + # Need to adjust coordinates to match grid_sample with align_corners=False + # grid_sample maps [-1, 1] to [0, W-1] when align_corners=False + # So our [0, W-1] input should be treated as [0.5, W-0.5] in pixel space + x_normalized = x + 0.5 + + if mode == 'nearest': + # Nearest neighbor sampling with zero padding outside [0, W-1] + if align_corners: + x_nearest = torch.round(x_normalized).long() + else: + x_nearest = torch.floor(x_normalized).long() + valid = (x_nearest >= 0) & (x_nearest < W) # (B, W_out) + x_index = torch.clamp(x_nearest, 0, W-1) + sampled = torch.gather(img, 3, x_index.view(B,1,1,-1).expand(B,C,1,-1)) + sampled = sampled * valid.view(B,1,1,-1).to(img.dtype) + + else: # bilinear + # Get integer and fractional parts + x_floor = torch.floor(x_normalized) + x_ceil = x_floor + 1 + x_frac = x_normalized - x_floor # (B, W_out) + + # Zero padding behavior: mark validity and zero-out invalid contributions + valid_floor = (x_floor >= 0) & (x_floor < W) + valid_ceil = (x_ceil >= 0) & (x_ceil < W) + x_floor_clamped = torch.clamp(x_floor, 0, W-1) + x_ceil_clamped = torch.clamp(x_ceil, 0, W-1) + + # Create index tensors + batch_idx = torch.arange(B, device=img.device).view(B, 1) + img_floor = torch.gather(img, 3, x_floor_clamped.view(B,1,1,-1).expand(B,C,1,-1).long()) + img_ceil = torch.gather(img, 3, x_ceil_clamped.view(B,1,1,-1).expand(B,C,1,-1).long()) + + # Apply validity masks (zero out-of-bounds samples) + img_floor = img_floor * valid_floor.view(B,1,1,-1).to(img.dtype) + img_ceil = img_ceil * valid_ceil.view(B,1,1,-1).to(img.dtype) + + # Linear interpolation + x_frac = x_frac.view(B,1,1,-1) + sampled = img_floor * (1 - x_frac) + img_ceil * x_frac + + return sampled + + +def bilinear_sampler(img, coords, mode='bilinear', mask=False, low_memory=False, use1d=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + coords[...,0] = 2*coords[...,0]/(W-1) - 1 + if low_memory: + B = img.shape[0] + out = [] + bs = 102400 + for b in np.arange(0,B,bs): + tmp = F.grid_sample(img[b:b+bs], coords[b:b+bs], align_corners=True) + out.append(tmp) + img = torch.cat(out, dim=0) + else: + img = F.grid_sample(img, coords, align_corners=True) + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + return img + diff --git a/cv/fast-foundationstereo/cpp/CMakeLists.txt b/cv/fast-foundationstereo/cpp/CMakeLists.txt new file mode 100644 index 0000000..107bf2d --- /dev/null +++ b/cv/fast-foundationstereo/cpp/CMakeLists.txt @@ -0,0 +1,156 @@ +cmake_minimum_required(VERSION 3.18) + +project(fast_foundation_stereo_cpp LANGUAGES CXX CUDA) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) + +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE "Release" CACHE STRING + "Build type: Debug, Release, RelWithDebInfo" FORCE) +endif() + +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES "80;86;89;90") +endif() + +add_compile_options( + $<$:--expt-relaxed-constexpr> + $<$:--diag-suppress=20015> + $<$:--diag-suppress=20013> + $<$:--diag-suppress=20011> + $<$:--diag-suppress=20091> +) + +find_package(CUDAToolkit REQUIRED) + +# ---- TensorRT ---- +set(TENSORRT_ROOT "/usr" CACHE PATH "TensorRT installation path") +set(TENSORRT_HINT_DIRS + ${TENSORRT_ROOT} + $ENV{CONDA_PREFIX} + /usr + /usr/local + /usr/local/tensorrt + ${CUDAToolkit_TARGET_DIR} + /usr/local/cuda + /usr/local/cuda/targets/x86_64-linux + /usr/local/cuda-13.0 + /usr/local/cuda-13.0/targets/x86_64-linux + /usr/src/tensorrt +) +if(DEFINED ENV{CONDA_PREFIX}) + file(GLOB TENSORRT_PY_LIB_DIRS + "$ENV{CONDA_PREFIX}/lib/python*/site-packages/tensorrt_libs" + ) + file(GLOB TENSORRT_PY_INCLUDE_DIRS + "$ENV{CONDA_PREFIX}/lib/python*/site-packages/tensorrt/include" + ) +endif() + +find_library(NVINFER_LIB + NAMES nvinfer nvinfer.so.10 + HINTS + ${TENSORRT_HINT_DIRS} + ${TENSORRT_PY_LIB_DIRS} + PATH_SUFFIXES + lib + lib64 + lib/x86_64-linux-gnu +) +if(NOT NVINFER_LIB) + message(FATAL_ERROR "TensorRT nvinfer library not found. Set -DTENSORRT_ROOT=/path/to/tensorrt") +endif() +message(STATUS "TensorRT nvinfer: ${NVINFER_LIB}") + +find_library(NVONNXPARSER_LIB + NAMES nvonnxparser nvonnxparser.so.10 + HINTS + ${TENSORRT_HINT_DIRS} + ${TENSORRT_PY_LIB_DIRS} + PATH_SUFFIXES + lib + lib64 + lib/x86_64-linux-gnu +) +if(NOT NVONNXPARSER_LIB) + message(FATAL_ERROR "TensorRT nvonnxparser library not found. Set -DTENSORRT_ROOT=/path/to/tensorrt") +endif() +message(STATUS "TensorRT nvonnxparser: ${NVONNXPARSER_LIB}") + +find_path(NVINFER_INCLUDE_DIR NvInfer.h + HINTS + ${TENSORRT_HINT_DIRS} + ${TENSORRT_PY_INCLUDE_DIRS} + PATH_SUFFIXES + include + include/x86_64-linux-gnu +) +if(NOT NVINFER_INCLUDE_DIR) + message(FATAL_ERROR "TensorRT headers (NvInfer.h) not found. Set -DTENSORRT_ROOT=/path/to/tensorrt") +endif() + +find_package(OpenCV QUIET COMPONENTS core imgcodecs imgproc) + +add_library(ffs_gwc_plugin SHARED + src/gwc_volume_plugin.cpp + src/depth_kernels.cu +) +set_target_properties(ffs_gwc_plugin PROPERTIES + CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON + OUTPUT_NAME "ffs_gwc_plugin" +) +target_include_directories(ffs_gwc_plugin PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${NVINFER_INCLUDE_DIR} +) +target_link_libraries(ffs_gwc_plugin PUBLIC + ${NVINFER_LIB} + CUDA::cudart +) + +add_library(ffs_depth_inference STATIC + src/ffs_depth_tensorrt.cpp + src/ffs_depth_single_tensorrt.cpp + src/gwc_volume_plugin.cpp + src/depth_kernels.cu +) +set_target_properties(ffs_depth_inference PROPERTIES + CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON +) +target_include_directories(ffs_depth_inference PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${NVINFER_INCLUDE_DIR} +) +target_link_libraries(ffs_depth_inference PUBLIC + ${NVINFER_LIB} + CUDA::cudart +) + +if(OpenCV_FOUND) + add_executable(ffs_depth_main app/main.cpp) + target_link_libraries(ffs_depth_main PRIVATE + ffs_depth_inference + ${OpenCV_LIBS} + ) +else() + message(WARNING "OpenCV C++ development package not found; skipping ffs_depth_main and ffs_profile_speed") +endif() + +add_executable(ffs_build_single_engine app/build_single_engine.cpp) +target_link_libraries(ffs_build_single_engine PRIVATE + ffs_depth_inference + ${NVONNXPARSER_LIB} +) + +if(OpenCV_FOUND) + add_executable(ffs_profile_speed app/profile_speed.cpp) + target_link_libraries(ffs_profile_speed PRIVATE + ffs_depth_inference + ${OpenCV_LIBS} + ) +endif() diff --git a/cv/fast-foundationstereo/cpp/README.md b/cv/fast-foundationstereo/cpp/README.md new file mode 100644 index 0000000..3afc133 --- /dev/null +++ b/cv/fast-foundationstereo/cpp/README.md @@ -0,0 +1,249 @@ +# Fast-FoundationStereo C++ Inference + +C++ runtime for Fast-FoundationStereo stereo depth inference on TensorRT. + +## Run from inside the C++ Docker container + +**All commands in this README are intended to be run from inside the container built from [`docker/dockerfile_cpp`](../docker/dockerfile_cpp).** That image carries the CUDA toolkit, TensorRT 10 runtime + ONNX parser + C++ headers, OpenCV development package, `trtexec`, and a Python environment with PyTorch / ONNX export tooling. The plain `docker/dockerfile` image does **not** ship the C++ TRT headers and will fail to build the C++ targets. + +Environment setup: + +```bash +docker build --network host -t ffs -f docker/dockerfile_cpp . +bash docker/run_container.sh +``` + + +`run_container.sh` mounts the parent of `docker/` (the repo root) into `/workspace`, so `cd /workspace//cpp` lands you in this folder. Every `cd cpp`, `cpp/build/...`, `python3 scripts/...`, and `trtexec ...` example below assumes that working directory. + +```text +cpp/ ++-- CMakeLists.txt ++-- README.md ++-- app/ +| +-- main.cpp -- ffs_depth_main: single image inference + visualization +| +-- build_single_engine.cpp -- ffs_build_single_engine: ONNX (with FFSGWCVolume plugin) -> .engine +| +-- profile_speed.cpp -- ffs_profile_speed: latency profiler for either route ++-- include/ +| +-- ffs_depth_tensorrt.hpp +| +-- ffs_depth_single_tensorrt.hpp +| +-- ffs_gwc_plugin.hpp ++-- src/ + +-- ffs_depth_tensorrt.cpp + +-- ffs_depth_single_tensorrt.cpp + +-- gwc_volume_plugin.cpp + +-- depth_kernels.cu +``` + +Two interchangeable inference routes are supported: + +- `FFSSingleEngineInference` loads **one** TensorRT engine that contains the `FFSGWCVolume` plugin node. +- `FFSDepthInference` loads **two** TensorRT engines (feature_runner + post_runner) and computes the GWC cost volume between them with a hand-written CUDA kernel. + +`ffs_depth_main` auto-detects the route based on the engine directory contents: + +- If the directory contains `fast_foundationstereo.engine`, it uses the single-engine plugin path. +- Otherwise it uses the two-engine reference path. + +Inputs and outputs are the same in both cases: + +- Input: a stereo image pair plus an intrinsic file (`demo_data/K.txt` format: 9 floats on line 1 for the 3x3 camera matrix, one float on line 2 for the stereo baseline in meters). +- Output: float32 disparity (input-pixel units), float32 depth in meters, and PNG visualizations. + +## Dependencies + +All C++ build dependencies are provided by the [`docker/dockerfile_cpp`](../docker/dockerfile_cpp) image (see [Run from inside the C++ Docker container](#run-from-inside-the-c-docker-container) for how to build and enter it): + +- CUDA Toolkit (for `nvcc` and the CUDA runtime). +- TensorRT 10 runtime, ONNX parser, and C++ headers. +- OpenCV development package (image I/O and depth visualization in `app/main.cpp`). +- `trtexec`, used by Route B to build the two-engine TensorRT engines from ONNX. + +## Build + +```bash +cd cpp +cmake -B build +cmake --build build -j +``` + +This produces: + +- `build/libffs_gwc_plugin.so` -- the FFSGWCVolume plugin as a shared library (loadable by `trtexec --staticPlugins=...`). +- `build/libffs_depth_inference.a` -- the inference static library. +- `build/ffs_build_single_engine` -- single-engine builder. +- `build/ffs_depth_main` -- demo / inference CLI. +- `build/ffs_profile_speed` -- latency profiler. + +## Route A: Single Engine with FFSGWCVolume Plugin + +### A.1 Export ONNX with the plugin node + +`scripts/make_plugin_onnx.py` exports one ONNX graph in which the GWC cost volume is represented by an `FFSGWCVolume` custom plugin node (resolved at engine-build time by `libffs_gwc_plugin.so`): + +```bash +python3 scripts/make_plugin_onnx.py \ + --model_dir weights/23-36-37/model_best_bp2_serialize.pth \ + --save_path output_plugin_onnx \ + --height 480 \ + --width 640 \ + --valid_iters 8 \ + --max_disp 192 +``` + +This writes: + +```text +output_plugin_onnx/ ++-- fast_foundationstereo_plugin.onnx ++-- onnx.yaml +``` + +### A.2 Build the single TensorRT engine + +```bash +cpp/build/ffs_build_single_engine \ + output_plugin_onnx/fast_foundationstereo_plugin.onnx \ + output_plugin_onnx/fast_foundationstereo.engine +``` + +By default the engine is built with FP16 enabled. Pass `--fp32` to disable FP16. Pass `--workspace-mb N` to override the workspace (default 4096 MB). + +After this step `output_plugin_onnx/` contains everything the runtime needs: + +```text +output_plugin_onnx/ ++-- fast_foundationstereo.engine ++-- onnx.yaml +``` + +### A.3 Run inference + +Run from the repository root: + +```bash +cpp/build/ffs_depth_main \ + output_plugin_onnx \ + demo_data/left.png \ + demo_data/right.png \ + demo_data/K.txt \ + output_plugin_onnx +``` + +The last argument is the output directory (default `ffs_output`). `ffs_depth_main` sees `fast_foundationstereo.engine` inside `output_plugin_onnx/` and uses `FFSSingleEngineInference`, which deserializes the engine and registers the `FFSGWCVolume` plugin before inference. + +### A.4 Python alternative (build engine and run inference without the C++ apps) + +The same plugin ONNX can be turned into an engine and executed end-to-end from Python. Only the C++ plugin shared library (`libffs_gwc_plugin.so`) is required from the C++ build; the C++ apps (`ffs_build_single_engine`, `ffs_depth_main`) are not. + +```bash +# 1. Export plugin ONNX +python3 scripts/make_plugin_onnx.py \ + --model_dir weights/23-36-37/model_best_bp2_serialize.pth \ + --save_path output_plugin_onnx \ + --height 480 \ + --width 640 + +# 2. Build the C++ plugin shared library +cmake -S cpp -B cpp/build +cmake --build cpp/build -j + +# 3. Build the TensorRT engine from Python +python3 scripts/build_plugin_trt.py \ + output_plugin_onnx/fast_foundationstereo_plugin.onnx \ + output_plugin_onnx/fast_foundationstereo.engine + +# 4. Run inference from Python +python3 scripts/run_demo_plugin_trt.py \ + --model_dir output_plugin_onnx \ + --left_file demo_data/left.png \ + --right_file demo_data/right.png \ + --intrinsic_file demo_data/K.txt \ + --out_dir output_plugin_onnx +``` + +Both `build_plugin_trt.py` and `run_demo_plugin_trt.py` auto-discover `libffs_gwc_plugin.so` in `cpp/build/`. Pass `--plugin_lib /path/to/libffs_gwc_plugin.so` to override the location. `build_plugin_trt.py` accepts `--fp32` and `--workspace-mb N` with the same meaning as `ffs_build_single_engine`. `run_demo_plugin_trt.py` writes the same five output files as the C++ demo (see [Outputs](#outputs)). + +## Route B: Two TensorRT Engines (feature_runner + post_runner) + +### B.1 Build the two engines + +See the [Two-stage ONNX section in the top-level README](../readme.md#two-stage-onnx) for how to export `feature_runner.onnx` / `post_runner.onnx` with `scripts/make_onnx.py` and then build `feature_runner.engine` / `post_runner.engine` with `trtexec`. After running those steps your engine directory should contain: + +```text +output_two_onnx/ ++-- feature_runner.engine ++-- post_runner.engine ++-- onnx.yaml +``` + +The two engines do **not** use the FFSGWCVolume plugin: the GWC cost volume is computed externally on GPU by `cpp/src/depth_kernels.cu` between the two engine calls, so a plain `trtexec --fp16` build with no custom plugin library is enough. + +### B.2 Run inference + +```bash +cpp/build/ffs_depth_main \ + output_two_onnx \ + demo_data/left.png \ + demo_data/right.png \ + demo_data/K.txt \ + output_two_onnx +``` + +`ffs_depth_main` sees no `fast_foundationstereo.engine` in the directory and falls back to `FFSDepthInference`, which executes the feature engine, builds the GWC volume with the CUDA kernel, then executes the post engine. + +## Outputs + +Both routes write the same files into the output directory: + +- `disparity.bin` -- raw float32 disparity (input-pixel units), prefixed by int32 `[height, width]`. +- `depth_meter.bin` -- raw float32 depth in meters, prefixed by int32 `[height, width]`. +- `depth_meter.npy` -- NumPy float32 depth in meters, shape `[height, width]`. +- `disp_vis.png` -- left/right/colorized-disparity side-by-side visualization. +- `depth_vis.png` -- colorized depth visualization. + +## Profile Speed + +`ffs_profile_speed` benchmarks either route end-to-end on a single image pair using CUDA events for GPU-side timing and `steady_clock` for host wall time. It auto-detects the route the same way `ffs_depth_main` does (presence of `fast_foundationstereo.engine` in the engine directory selects single-engine), and `--mode` can be set explicitly: + +```bash +cpp/build/ffs_profile_speed \ + [--mode auto|two|single] [--warmup N] [--runs N] [--include-depth] +``` + +Defaults: `--mode auto`, `--warmup 10`, `--runs 30`. With `--include-depth`, the disparity-to-depth conversion is included in the timed region (otherwise only the `infer()` call is timed). + +### Profile the single engine + +```bash +cpp/build/ffs_profile_speed \ + output_plugin_onnx \ + demo_data/left.png demo_data/right.png demo_data/K.txt \ + --mode single --warmup 20 --runs 100 +``` + +### Profile the two engines + +```bash +cpp/build/ffs_profile_speed \ + output_two_onnx \ + demo_data/left.png demo_data/right.png demo_data/K.txt \ + --mode two --warmup 20 --runs 100 +``` + +The output looks like: + +```text +mode=single +image=960x540 +model=640x480 +warmup=20 runs=100 +timed_region=infer +gpu mean_ms=... p50_ms=... p90_ms=... min_ms=... max_ms=... std_ms=... +host mean_ms=... p50_ms=... p90_ms=... min_ms=... max_ms=... std_ms=... +``` + +- `gpu` is `cudaEventElapsedTime` between start/stop events on the inference stream (pure GPU work). +- `host` is `std::chrono::steady_clock` around the same region, including the cost of `cudaEventSynchronize`. `host >= gpu` always. + +For a side-by-side comparison, run the profiler twice with different engine directories and `--mode`, then compare the `gpu mean_ms` columns. diff --git a/cv/fast-foundationstereo/cpp/app/build_single_engine.cpp b/cv/fast-foundationstereo/cpp/app/build_single_engine.cpp new file mode 100644 index 0000000..cc0bd3c --- /dev/null +++ b/cv/fast-foundationstereo/cpp/app/build_single_engine.cpp @@ -0,0 +1,120 @@ +#include "ffs_gwc_plugin.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +class Logger : public nvinfer1::ILogger { +public: + void log(Severity severity, const char* msg) noexcept override { + if (severity <= Severity::kINFO) { + std::cerr << "[TRT] " << msg << "\n"; + } + } +}; + +struct TrtDestroy { + template + void operator()(T* p) const { + delete p; + } +}; + +template +using TrtPtr = std::unique_ptr; + +void printUsage(const char* prog) { + std::cerr + << "Usage: " << prog << " [--fp32] [--workspace-mb N]\n" + << "\n" + << "Builds a single TensorRT engine from an ONNX graph containing the\n" + << "FFSGWCVolume custom plugin node.\n"; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 3) { + printUsage(argv[0]); + return 1; + } + + try { + const std::filesystem::path onnx_path = argv[1]; + const std::filesystem::path engine_path = argv[2]; + bool fp16 = true; + size_t workspace_mb = 4096; + + for (int i = 3; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--fp32") { + fp16 = false; + } else if (arg == "--workspace-mb" && i + 1 < argc) { + workspace_mb = static_cast(std::stoull(argv[++i])); + } else { + throw std::runtime_error("unknown argument: " + arg); + } + } + + if (!std::filesystem::exists(onnx_path)) { + throw std::runtime_error("ONNX file does not exist: " + onnx_path.string()); + } + if (!engine_path.parent_path().empty()) { + std::filesystem::create_directories(engine_path.parent_path()); + } + + if (!ffs_depth::registerFFSGWCPlugin()) { + throw std::runtime_error("failed to register FFSGWCVolume plugin"); + } + + Logger logger; + TrtPtr builder(nvinfer1::createInferBuilder(logger)); + if (!builder) throw std::runtime_error("createInferBuilder failed"); + + const auto explicit_batch = + 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); + TrtPtr network(builder->createNetworkV2(explicit_batch)); + if (!network) throw std::runtime_error("createNetworkV2 failed"); + + TrtPtr parser(nvonnxparser::createParser(*network, logger)); + if (!parser) throw std::runtime_error("createParser failed"); + + if (!parser->parseFromFile(onnx_path.string().c_str(), + static_cast(nvinfer1::ILogger::Severity::kWARNING))) { + for (int32_t i = 0; i < parser->getNbErrors(); ++i) { + std::cerr << parser->getError(i)->desc() << "\n"; + } + throw std::runtime_error("failed to parse ONNX: " + onnx_path.string()); + } + + TrtPtr config(builder->createBuilderConfig()); + if (!config) throw std::runtime_error("createBuilderConfig failed"); + config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, + workspace_mb * 1024ULL * 1024ULL); + if (fp16 && builder->platformHasFastFp16()) { + config->setFlag(nvinfer1::BuilderFlag::kFP16); + } + + TrtPtr serialized(builder->buildSerializedNetwork(*network, *config)); + if (!serialized) throw std::runtime_error("buildSerializedNetwork failed"); + + std::ofstream out(engine_path, std::ios::binary); + if (!out) throw std::runtime_error("cannot write engine: " + engine_path.string()); + out.write(static_cast(serialized->data()), serialized->size()); + + std::cout << "Built engine: " << engine_path << "\n"; + std::cout << "Precision: " << (fp16 ? "FP16 allowed" : "FP32") << "\n"; + } catch (const std::exception& e) { + std::cerr << "ERROR: " << e.what() << "\n"; + return 1; + } + return 0; +} diff --git a/cv/fast-foundationstereo/cpp/app/main.cpp b/cv/fast-foundationstereo/cpp/app/main.cpp new file mode 100644 index 0000000..b21968a --- /dev/null +++ b/cv/fast-foundationstereo/cpp/app/main.cpp @@ -0,0 +1,346 @@ +/** + * Standalone Fast-FoundationStereo TensorRT inference demo. + * + * Usage: + * ./ffs_depth_main [output_dir] + * + * Outputs: + * /disparity.bin int32 H, int32 W, then H*W float32 disparity + * /depth_meter.bin int32 H, int32 W, then H*W float32 depth in meters + * /depth_meter.npy NumPy float32 depth in meters, shape (H, W) + * /disp_vis.png left/right/colorized-disparity visualization + * /depth_vis.png colorized depth visualization + */ + +#include "ffs_depth_tensorrt.hpp" +#include "ffs_depth_single_tensorrt.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +void checkCuda(cudaError_t status, const char* what) { + if (status != cudaSuccess) { + throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status)); + } +} + +void printUsage(const char* prog) { + std::cerr + << "Usage: " << prog << " [output_dir]\n" + << "\n" + << " engine_dir : directory containing feature_runner.engine, post_runner.engine, onnx.yaml\n" + << " left_image : left stereo image readable by OpenCV\n" + << " right_image: right stereo image with the same size as left_image\n" + << " intrinsic_file: text file with 3x3 K on line 1 and baseline in meters on line 2\n" + << " output_dir : output directory (default: ffs_output)\n"; +} + +struct CudaBuffer { + void* ptr = nullptr; + + ~CudaBuffer() { + if (ptr) cudaFree(ptr); + } + + CudaBuffer() = default; + CudaBuffer(const CudaBuffer&) = delete; + CudaBuffer& operator=(const CudaBuffer&) = delete; + + void allocate(size_t bytes) { + checkCuda(cudaMalloc(&ptr, bytes), "cudaMalloc"); + } + + template + T* as() { + return static_cast(ptr); + } +}; + +struct Intrinsics { + float k[9] = {}; + float baseline = 0.0f; +}; + +Intrinsics loadIntrinsics(const std::string& path) { + std::ifstream in(path); + if (!in) { + throw std::runtime_error("cannot open intrinsic file: " + path); + } + + std::string k_line; + std::string baseline_line; + if (!std::getline(in, k_line) || !std::getline(in, baseline_line)) { + throw std::runtime_error("intrinsic file must contain K on line 1 and baseline on line 2"); + } + + Intrinsics intr; + std::istringstream k_stream(k_line); + for (float& value : intr.k) { + if (!(k_stream >> value)) { + throw std::runtime_error("intrinsic file K line must contain 9 float values"); + } + } + + std::istringstream baseline_stream(baseline_line); + if (!(baseline_stream >> intr.baseline)) { + throw std::runtime_error("intrinsic file baseline line must contain one float value"); + } + if (!(intr.k[0] > 0.0f) || !(intr.baseline > 0.0f)) { + throw std::runtime_error("invalid focal length or baseline in intrinsic file"); + } + + return intr; +} + +void saveFloatMatrix(const std::filesystem::path& path, + const std::vector& values, + int height, + int width) { + std::ofstream out(path, std::ios::binary); + if (!out) { + throw std::runtime_error("cannot write matrix: " + path.string()); + } + const int32_t dims[2] = {height, width}; + out.write(reinterpret_cast(dims), sizeof(dims)); + out.write(reinterpret_cast(values.data()), + values.size() * sizeof(float)); +} + +void saveNpyFloat32(const std::filesystem::path& path, + const std::vector& values, + int height, + int width) { + std::ofstream out(path, std::ios::binary); + if (!out) { + throw std::runtime_error("cannot write npy matrix: " + path.string()); + } + + std::ostringstream header_stream; + header_stream << "{'descr': ' 65535) { + throw std::runtime_error("npy header too large: " + path.string()); + } + const uint16_t header_len = static_cast(header.size()); + + out.write("\x93NUMPY", 6); + const char version[2] = {1, 0}; + out.write(version, 2); + out.write(reinterpret_cast(&header_len), sizeof(header_len)); + out.write(header.data(), static_cast(header.size())); + out.write(reinterpret_cast(values.data()), + static_cast(values.size() * sizeof(float))); +} + +cv::Mat colorizeDisparity(const std::vector& disparity, int height, int width) { + std::vector safe_disp(disparity.size(), 0.0f); + std::vector valid_mask(disparity.size(), 0); + for (size_t i = 0; i < disparity.size(); ++i) { + if (std::isfinite(disparity[i])) { + safe_disp[i] = std::max(disparity[i], 0.0f); + valid_mask[i] = 255; + } + } + + cv::Mat disp_mat(height, width, CV_32FC1, safe_disp.data()); + cv::Mat valid(height, width, CV_8UC1, valid_mask.data()); + + if (cv::countNonZero(valid) == 0) { + return cv::Mat::zeros(height, width, CV_8UC3); + } + + double min_val = 0.0; + double max_val = 0.0; + cv::minMaxLoc(disp_mat, &min_val, &max_val, nullptr, nullptr, valid); + if (max_val <= min_val) { + max_val = min_val + 1.0; + } + + cv::Mat disp_u8; + disp_mat.convertTo(disp_u8, CV_8UC1, 255.0 / (max_val - min_val), + -255.0 * min_val / (max_val - min_val)); + + cv::Mat colored; + cv::applyColorMap(disp_u8, colored, cv::COLORMAP_TURBO); + colored.setTo(cv::Scalar(0, 0, 0), ~valid); + return colored; +} + +cv::Mat colorizeDepth(const std::vector& depth, int height, int width) { + std::vector safe_depth(depth.size(), 0.0f); + std::vector valid_mask(depth.size(), 0); + for (size_t i = 0; i < depth.size(); ++i) { + if (std::isfinite(depth[i]) && depth[i] > 0.0f) { + safe_depth[i] = depth[i]; + valid_mask[i] = 255; + } + } + + cv::Mat depth_mat(height, width, CV_32FC1, safe_depth.data()); + cv::Mat valid(height, width, CV_8UC1, valid_mask.data()); + + if (cv::countNonZero(valid) == 0) { + return cv::Mat::zeros(height, width, CV_8UC3); + } + + double min_val = 0.0; + double max_val = 0.0; + cv::minMaxLoc(depth_mat, &min_val, &max_val, nullptr, nullptr, valid); + if (max_val <= min_val) { + max_val = min_val + 1.0; + } + + cv::Mat depth_u8; + depth_mat.convertTo(depth_u8, CV_8UC1, 255.0 / (max_val - min_val), + -255.0 * min_val / (max_val - min_val)); + + cv::Mat colored; + cv::applyColorMap(depth_u8, colored, cv::COLORMAP_TURBO); + colored.setTo(cv::Scalar(0, 0, 0), ~valid); + return colored; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 5 || argc > 6) { + printUsage(argv[0]); + return 1; + } + + try { + const std::string engine_dir = argv[1]; + const std::string left_path = argv[2]; + const std::string right_path = argv[3]; + const std::string intrinsic_path = argv[4]; + const std::filesystem::path output_dir = (argc == 6) ? argv[5] : "ffs_output"; + const Intrinsics intr = loadIntrinsics(intrinsic_path); + + cv::Mat left_bgr = cv::imread(left_path, cv::IMREAD_COLOR); + cv::Mat right_bgr = cv::imread(right_path, cv::IMREAD_COLOR); + + if (left_bgr.empty()) { + throw std::runtime_error("cannot read left image: " + left_path); + } + if (right_bgr.empty()) { + throw std::runtime_error("cannot read right image: " + right_path); + } + if (left_bgr.size() != right_bgr.size()) { + throw std::runtime_error("left and right images must have identical dimensions"); + } + if (!left_bgr.isContinuous()) left_bgr = left_bgr.clone(); + if (!right_bgr.isContinuous()) right_bgr = right_bgr.clone(); + + std::filesystem::create_directories(output_dir); + + const int height = left_bgr.rows; + const int width = left_bgr.cols; + const size_t image_bytes = static_cast(height) * width * 3 * sizeof(uint8_t); + const size_t map_bytes = static_cast(height) * width * sizeof(float); + + CudaBuffer d_left; + CudaBuffer d_right; + CudaBuffer d_disparity; + CudaBuffer d_depth; + d_left.allocate(image_bytes); + d_right.allocate(image_bytes); + d_disparity.allocate(map_bytes); + d_depth.allocate(map_bytes); + + checkCuda(cudaMemcpy(d_left.ptr, left_bgr.data, image_bytes, cudaMemcpyHostToDevice), + "cudaMemcpy left image"); + checkCuda(cudaMemcpy(d_right.ptr, right_bgr.data, image_bytes, cudaMemcpyHostToDevice), + "cudaMemcpy right image"); + + std::cout << "Input images: " << width << "x" << height << "\n"; + std::cout << "Depth: fx=" << intr.k[0] << " baseline=" << intr.baseline << " m\n"; + std::cout << "Loading engines from: " << engine_dir << "\n"; + const bool use_single_engine = + std::filesystem::exists(std::filesystem::path(engine_dir) / "fast_foundationstereo.engine"); + + if (use_single_engine) { + std::cout << "Runtime: single TensorRT engine with FFSGWCVolume plugin\n"; + ffs_depth::FFSSingleEngineInference ffs(engine_dir); + ffs.infer(d_left.as(), d_right.as(), height, width, + d_disparity.as()); + ffs.dispToDepth(d_disparity.as(), height, width, + intr.k[0], intr.baseline, d_depth.as()); + ffs.sync(); + } else { + std::cout << "Runtime: two TensorRT engines with external CUDA GWC\n"; + ffs_depth::FFSDepthInference ffs(engine_dir); + ffs.infer(d_left.as(), d_right.as(), height, width, + d_disparity.as()); + ffs.dispToDepth(d_disparity.as(), height, width, + intr.k[0], intr.baseline, d_depth.as()); + ffs.sync(); + } + + std::vector disparity(static_cast(height) * width); + std::vector depth(static_cast(height) * width); + checkCuda(cudaMemcpy(disparity.data(), d_disparity.ptr, map_bytes, cudaMemcpyDeviceToHost), + "cudaMemcpy disparity"); + checkCuda(cudaMemcpy(depth.data(), d_depth.ptr, map_bytes, cudaMemcpyDeviceToHost), + "cudaMemcpy depth"); + + const auto disparity_path = output_dir / "disparity.bin"; + saveFloatMatrix(disparity_path, disparity, height, width); + + const auto depth_path = output_dir / "depth_meter.bin"; + saveFloatMatrix(depth_path, depth, height, width); + + const auto depth_npy_path = output_dir / "depth_meter.npy"; + saveNpyFloat32(depth_npy_path, depth, height, width); + + const cv::Mat disp_color = colorizeDisparity(disparity, height, width); + cv::Mat disp_vis; + cv::hconcat(std::vector{left_bgr, right_bgr, disp_color}, disp_vis); + const auto disp_vis_path = output_dir / "disp_vis.png"; + if (!cv::imwrite(disp_vis_path.string(), disp_vis)) { + throw std::runtime_error("failed to write disparity visualization: " + disp_vis_path.string()); + } + + const cv::Mat depth_vis = colorizeDepth(depth, height, width); + const auto depth_vis_path = output_dir / "depth_vis.png"; + if (!cv::imwrite(depth_vis_path.string(), depth_vis)) { + throw std::runtime_error("failed to write depth visualization: " + depth_vis_path.string()); + } + + std::cout << "Saved: " << disparity_path << "\n"; + std::cout << "Saved: " << depth_path << "\n"; + std::cout << "Saved: " << depth_npy_path << "\n"; + std::cout << "Saved: " << disp_vis_path << "\n"; + std::cout << "Saved: " << depth_vis_path << "\n"; + } catch (const std::exception& e) { + std::cerr << "ERROR: " << e.what() << "\n"; + return 1; + } + + return 0; +} diff --git a/cv/fast-foundationstereo/cpp/app/profile_speed.cpp b/cv/fast-foundationstereo/cpp/app/profile_speed.cpp new file mode 100644 index 0000000..61e29bf --- /dev/null +++ b/cv/fast-foundationstereo/cpp/app/profile_speed.cpp @@ -0,0 +1,312 @@ +#include "ffs_depth_tensorrt.hpp" +#include "ffs_depth_single_tensorrt.hpp" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +void checkCuda(cudaError_t status, const char* what) { + if (status != cudaSuccess) { + throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status)); + } +} + +struct CudaBuffer { + void* ptr = nullptr; + ~CudaBuffer() { if (ptr) cudaFree(ptr); } + void allocate(size_t bytes) { checkCuda(cudaMalloc(&ptr, bytes), "cudaMalloc"); } + template T* as() { return static_cast(ptr); } +}; + +struct Intrinsics { + float k[9] = {}; + float baseline = 0.0f; +}; + +struct Args { + std::string engine_dir; + std::string left_path; + std::string right_path; + std::string intrinsic_path; + std::string mode = "auto"; + int warmup = 10; + int runs = 30; + bool include_depth = false; +}; + +struct Stats { + double mean = 0.0; + double min = 0.0; + double p50 = 0.0; + double p90 = 0.0; + double max = 0.0; + double stddev = 0.0; +}; + +void printUsage(const char* prog) { + std::cerr + << "Usage: " << prog + << " " + << " [--mode auto|two|single] [--warmup N] [--runs N] [--include-depth]\n"; +} + +Args parseArgs(int argc, char** argv) { + if (argc < 5) { + printUsage(argv[0]); + throw std::runtime_error("missing required arguments"); + } + + Args args; + args.engine_dir = argv[1]; + args.left_path = argv[2]; + args.right_path = argv[3]; + args.intrinsic_path = argv[4]; + + for (int i = 5; i < argc; ++i) { + const std::string key = argv[i]; + auto requireValue = [&](const char* name) -> std::string { + if (i + 1 >= argc) { + throw std::runtime_error(std::string("missing value for ") + name); + } + return argv[++i]; + }; + + if (key == "--mode") { + args.mode = requireValue("--mode"); + } else if (key == "--warmup") { + args.warmup = std::stoi(requireValue("--warmup")); + } else if (key == "--runs") { + args.runs = std::stoi(requireValue("--runs")); + } else if (key == "--include-depth") { + args.include_depth = true; + } else { + throw std::runtime_error("unknown argument: " + key); + } + } + + if (args.mode != "auto" && args.mode != "two" && args.mode != "single") { + throw std::runtime_error("--mode must be auto, two, or single"); + } + if (args.warmup < 0 || args.runs <= 0) { + throw std::runtime_error("--warmup must be >= 0 and --runs must be > 0"); + } + return args; +} + +Intrinsics loadIntrinsics(const std::string& path) { + std::ifstream in(path); + if (!in) throw std::runtime_error("cannot open intrinsic file: " + path); + std::string k_line; + std::string baseline_line; + if (!std::getline(in, k_line) || !std::getline(in, baseline_line)) { + throw std::runtime_error("intrinsic file must contain K on line 1 and baseline on line 2"); + } + + Intrinsics intr; + std::istringstream ks(k_line); + for (float& v : intr.k) { + if (!(ks >> v)) throw std::runtime_error("K line must contain 9 floats"); + } + std::istringstream bs(baseline_line); + if (!(bs >> intr.baseline)) throw std::runtime_error("baseline line must contain one float"); + return intr; +} + +Stats summarize(std::vector values) { + if (values.empty()) throw std::runtime_error("cannot summarize empty timing vector"); + std::sort(values.begin(), values.end()); + const double sum = std::accumulate(values.begin(), values.end(), 0.0); + const double mean = sum / static_cast(values.size()); + double var = 0.0; + for (float v : values) { + const double d = static_cast(v) - mean; + var += d * d; + } + var /= static_cast(values.size()); + + auto percentile = [&](double p) { + const size_t idx = static_cast( + std::llround((values.size() - 1) * p / 100.0)); + return static_cast(values[std::min(idx, values.size() - 1)]); + }; + + Stats stats; + stats.mean = mean; + stats.min = values.front(); + stats.p50 = percentile(50.0); + stats.p90 = percentile(90.0); + stats.max = values.back(); + stats.stddev = std::sqrt(var); + return stats; +} + +void printStats(const char* label, const Stats& s) { + std::cout + << label + << " mean_ms=" << s.mean + << " p50_ms=" << s.p50 + << " p90_ms=" << s.p90 + << " min_ms=" << s.min + << " max_ms=" << s.max + << " std_ms=" << s.stddev + << "\n"; +} + +template +std::vector profileRunner( + Runner& runner, + uint8_t* d_left, + uint8_t* d_right, + int height, + int width, + float* d_disp, + float* d_depth, + float fx, + float baseline, + int warmup, + int runs, + bool include_depth, + std::vector& host_ms) +{ + cudaEvent_t start = nullptr; + cudaEvent_t stop = nullptr; + checkCuda(cudaEventCreate(&start), "cudaEventCreate start"); + checkCuda(cudaEventCreate(&stop), "cudaEventCreate stop"); + + std::vector gpu_ms; + gpu_ms.reserve(static_cast(runs)); + host_ms.clear(); + host_ms.reserve(static_cast(runs)); + + for (int i = 0; i < warmup + runs; ++i) { + const auto host_start = std::chrono::steady_clock::now(); + checkCuda(cudaEventRecord(start, runner.stream()), "cudaEventRecord start"); + runner.infer(d_left, d_right, height, width, d_disp); + if (include_depth) { + runner.dispToDepth(d_disp, height, width, fx, baseline, d_depth); + } + checkCuda(cudaEventRecord(stop, runner.stream()), "cudaEventRecord stop"); + checkCuda(cudaEventSynchronize(stop), "cudaEventSynchronize stop"); + const auto host_stop = std::chrono::steady_clock::now(); + + if (i >= warmup) { + float elapsed = 0.0f; + checkCuda(cudaEventElapsedTime(&elapsed, start, stop), "cudaEventElapsedTime"); + gpu_ms.push_back(elapsed); + host_ms.push_back(static_cast( + std::chrono::duration(host_stop - host_start).count())); + } + } + + cudaEventDestroy(stop); + cudaEventDestroy(start); + return gpu_ms; +} + +template +void runProfile(const Args& args, + Runner& runner, + uint8_t* d_left, + uint8_t* d_right, + int height, + int width, + float* d_disp, + float* d_depth, + const Intrinsics& intr, + const char* mode_name) { + std::vector host_ms; + const std::vector gpu_ms = profileRunner( + runner, + d_left, + d_right, + height, + width, + d_disp, + d_depth, + intr.k[0], + intr.baseline, + args.warmup, + args.runs, + args.include_depth, + host_ms); + + std::cout << "mode=" << mode_name << "\n"; + std::cout << "image=" << width << "x" << height << "\n"; + std::cout << "model=" << runner.modelWidth() << "x" << runner.modelHeight() << "\n"; + std::cout << "warmup=" << args.warmup << " runs=" << args.runs << "\n"; + std::cout << "timed_region=" << (args.include_depth ? "infer+dispToDepth" : "infer") << "\n"; + printStats("gpu", summarize(gpu_ms)); + printStats("host", summarize(host_ms)); +} + +} // namespace + +int main(int argc, char** argv) { + try { + const Args args = parseArgs(argc, argv); + const Intrinsics intr = loadIntrinsics(args.intrinsic_path); + + cv::Mat left = cv::imread(args.left_path, cv::IMREAD_COLOR); + cv::Mat right = cv::imread(args.right_path, cv::IMREAD_COLOR); + if (left.empty()) throw std::runtime_error("cannot read left image: " + args.left_path); + if (right.empty()) throw std::runtime_error("cannot read right image: " + args.right_path); + if (left.size() != right.size()) throw std::runtime_error("left/right size mismatch"); + if (!left.isContinuous()) left = left.clone(); + if (!right.isContinuous()) right = right.clone(); + + const int height = left.rows; + const int width = left.cols; + const size_t image_bytes = static_cast(height) * width * 3; + const size_t map_bytes = static_cast(height) * width * sizeof(float); + + CudaBuffer d_left; + CudaBuffer d_right; + CudaBuffer d_disp; + CudaBuffer d_depth; + d_left.allocate(image_bytes); + d_right.allocate(image_bytes); + d_disp.allocate(map_bytes); + d_depth.allocate(map_bytes); + checkCuda(cudaMemcpy(d_left.ptr, left.data, image_bytes, cudaMemcpyHostToDevice), "copy left"); + checkCuda(cudaMemcpy(d_right.ptr, right.data, image_bytes, cudaMemcpyHostToDevice), "copy right"); + + std::string mode = args.mode; + if (mode == "auto") { + mode = std::filesystem::exists( + std::filesystem::path(args.engine_dir) / "fast_foundationstereo.engine") + ? "single" + : "two"; + } + + if (mode == "single") { + ffs_depth::FFSSingleEngineInference runner(args.engine_dir); + runProfile(args, runner, d_left.as(), d_right.as(), + height, width, d_disp.as(), d_depth.as(), + intr, "single"); + } else { + ffs_depth::FFSDepthInference runner(args.engine_dir); + runProfile(args, runner, d_left.as(), d_right.as(), + height, width, d_disp.as(), d_depth.as(), + intr, "two"); + } + } catch (const std::exception& e) { + std::cerr << "ERROR: " << e.what() << "\n"; + return 1; + } + return 0; +} diff --git a/cv/fast-foundationstereo/cpp/include/ffs_depth_single_tensorrt.hpp b/cv/fast-foundationstereo/cpp/include/ffs_depth_single_tensorrt.hpp new file mode 100644 index 0000000..9f13eb3 --- /dev/null +++ b/cv/fast-foundationstereo/cpp/include/ffs_depth_single_tensorrt.hpp @@ -0,0 +1,80 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "ffs_depth_tensorrt.hpp" + +namespace ffs_depth { + +/** + * Single TensorRT engine inference path. + * + * Expected engine directory layout: + * - fast_foundationstereo.engine + * - onnx.yaml + * + * The engine is built from the plugin ONNX export path where the GWC cost + * volume is represented by an FFSGWCVolume TensorRT plugin node. + */ +class FFSSingleEngineInference { +public: + explicit FFSSingleEngineInference(const std::string& engine_dir); + ~FFSSingleEngineInference(); + + FFSSingleEngineInference(const FFSSingleEngineInference&) = delete; + FFSSingleEngineInference& operator=(const FFSSingleEngineInference&) = delete; + + void infer(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float* d_disp_out); + + void dispToDepth(const float* d_disp, + int height, int width, + float fx, float baseline_m, + float* d_depth_out); + + void inferDepth(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float fx, float baseline_m, + float* d_depth_out); + + void sync(); + + int modelHeight() const { return config_.image_height; } + int modelWidth() const { return config_.image_width; } + int maxDisp() const { return config_.max_disp; } + int cvGroup() const { return config_.cv_group; } + const FFSDepthInference::Config& config() const { return config_; } + cudaStream_t stream() const { return stream_; } + +private: + void loadConfig(const std::string& config_path); + void loadEngine(const std::string& path); + void allocateBuffers(); + void freeDeviceBuffers(); + void preprocessRGBGPU(const uint8_t* d_rgb, int src_h, int src_w, float* d_output); + + FFSDepthInference::Config config_; + + std::unique_ptr runtime_; + std::unique_ptr engine_; + std::unique_ptr context_; + + cudaStream_t stream_ = nullptr; + + float* d_left_ = nullptr; + float* d_right_ = nullptr; + float* d_disp_ = nullptr; + float* d_disp_cropped_ = nullptr; + float* d_disp_for_depth_ = nullptr; + + int64_t depth_alloc_pixels_ = 0; + int scaled_w_ = 0; + int scaled_h_ = 0; +}; + +} // namespace ffs_depth diff --git a/cv/fast-foundationstereo/cpp/include/ffs_depth_tensorrt.hpp b/cv/fast-foundationstereo/cpp/include/ffs_depth_tensorrt.hpp new file mode 100644 index 0000000..70be62c --- /dev/null +++ b/cv/fast-foundationstereo/cpp/include/ffs_depth_tensorrt.hpp @@ -0,0 +1,190 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace ffs_depth { + +/** + * FoundationStereo TensorRT depth inference. + * + * Uses the two-engine architecture (feature_runner + post_runner). + * The GWC (Group-wise Correlation) volume is computed on GPU between the two engines. + * + * Expected engine directory layout: + * - feature_runner.engine + * - post_runner.engine + * - onnx.yaml (contains image_size, max_disp, cv_group, valid_iters) + */ +class FFSDepthInference { +public: + struct Config { + int image_height = 480; // Model input height (must be divisible by 32) + int image_width = 864; // Model input width (must be divisible by 32) + int max_disp = 192; // Maximum disparity search range in pixels + int cv_group = 8; // Number of groups for group-wise correlation (GWC) volume + int valid_iters = 8; // Number of GRU refinement iterations during inference + bool normalize = true; // Whether to L2-normalize features in GWC correlation + }; + + /** + * Load engines and allocate all GPU buffers. + * @param engine_dir Directory containing the two .engine files and onnx.yaml. + */ + explicit FFSDepthInference(const std::string& engine_dir); + ~FFSDepthInference(); + + FFSDepthInference(const FFSDepthInference&) = delete; + FFSDepthInference& operator=(const FFSDepthInference&) = delete; + + /** + * Run stereo disparity inference entirely on the GPU. + * + * This method is asynchronous: all work (preprocessing, TensorRT + * inference, postprocessing) is enqueued on an internal CUDA stream + * and the call returns immediately. Call sync() to block until the + * output buffer is safe to read. + * + * Input images are expected in HWC uint8 BGR (OpenCV) format on device memory. + * If (input_h, input_w) differs from the model size, the images are + * uniformly (aspect-preserving) bilinearly resized down to fit the model + * input and the remaining right/bottom strip is filled with replicate + * padding. After inference the disparity is cropped back to the scaled + * region and nearest-neighbour upsampled to (input_h, input_w), with + * a horizontal scale correction so the values are in input-pixel units. + * + * @param d_left_rgb Left image on GPU (input_h x input_w x 3, uint8 BGR) + * @param d_right_rgb Right image on GPU (input_h x input_w x 3, uint8 BGR) + * @param input_h Height of input images (must be > 0) + * @param input_w Width of input images (must be > 0) + * @param d_disp_out Output disparity on GPU (input_h x input_w, float32, + * in input-pixel units, clamped to >= 0). + */ + void infer(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float* d_disp_out); + + /** + * Convert an input-resolution disparity map to float32 depth in meters. + * + * This method is asynchronous: work is enqueued on the internal CUDA + * stream and the call returns immediately. Call sync() before reading. + * + * Conversion semantics match scripts/run_demo.py (see inferDepth docstring + * for the +inf / off-image-correspondence handling). + * + * @param d_disp Input disparity on GPU (height x width, float32). + * Must be expressed in INPUT-image pixel units (which + * is what infer() produces). + * @param height Disparity / depth height + * @param width Disparity / depth width + * @param fx Focal length in pixels at INPUT resolution + * @param baseline_m Stereo baseline in meters + * @param d_depth_out Output depth on GPU (height x width, float32, meters) + */ + void dispToDepth(const float* d_disp, + int height, int width, + float fx, float baseline_m, + float* d_depth_out); + + /** + * Run stereo inference and convert disparity to float32 depth (meters) in one call. + * + * This method is asynchronous: all work is enqueued on an internal + * CUDA stream and the call returns immediately. Call sync() to block + * until d_depth_out is safe to read. + * + * The conversion matches scripts/run_demo.py: + * depth_m = fx * baseline_m / disparity + * where: + * - disparity is in INPUT-image pixel units (infer() already rescales it), + * so `fx` must come from the INPUT-resolution intrinsics (i.e. K[0,0] + * of the unscaled camera matrix, not the model-resolution intrinsics). + * - pixels with x - disparity < 0 (right-image correspondence off-image) + * are marked invalid: disparity is replaced with +inf, yielding depth 0. + * - disparity == 0 yields depth = +inf (consumers should mask non-finite + * values before using the depth map for downstream geometry). + * + * @param d_left_rgb Left image on GPU (input_h x input_w x 3, uint8 BGR/RGB) + * @param d_right_rgb Right image on GPU (input_h x input_w x 3, uint8 BGR/RGB) + * @param input_h Height of input images + * @param input_w Width of input images + * @param fx Focal length in pixels at INPUT resolution + * @param baseline_m Stereo baseline in meters + * @param d_depth_out Output depth on GPU (input_h x input_w, float32, meters) + */ + void inferDepth(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float fx, float baseline_m, + float* d_depth_out); + + /** Block until all async work on the internal CUDA stream has completed. */ + void sync(); + + int modelHeight() const { return config_.image_height; } + int modelWidth() const { return config_.image_width; } + int maxDisp() const { return config_.max_disp; } + int cvGroup() const { return config_.cv_group; } + const Config& config() const { return config_; } + cudaStream_t stream() const { return stream_; } + +private: + void loadConfig(const std::string& config_path); + void loadEngine(const std::string& path, + std::unique_ptr& engine, + std::unique_ptr& context); + void allocateBuffers(); + void allocateFeatureBuffers(); + void allocatePostBuffers(); + void freeDeviceBuffers(); // safe to call repeatedly; nulls every pointer + void preprocessRGBGPU(const uint8_t* d_rgb, int src_h, int src_w, float* d_output); + + void runFeatureRunner(); + void buildGWCVolume(); + void runPostRunner(); + + Config config_; + + std::unique_ptr runtime_; + std::unique_ptr feature_engine_; + std::unique_ptr feature_context_; + std::unique_ptr post_engine_; + std::unique_ptr post_context_; + + cudaStream_t stream_ = nullptr; + + float* d_left_ = nullptr; + float* d_right_ = nullptr; + + float* d_feat_left_04_ = nullptr; + float* d_feat_left_08_ = nullptr; + float* d_feat_left_16_ = nullptr; + float* d_feat_left_32_ = nullptr; + float* d_feat_right_04_ = nullptr; + float* d_stem_2x_ = nullptr; + + std::vector feat_04_dims_; + std::vector feat_08_dims_; + std::vector feat_16_dims_; + std::vector feat_32_dims_; + std::vector stem_2x_dims_; + + float* d_gwc_volume_ = nullptr; + int gwc_disp_levels_ = 0; + + float* d_disp_ = nullptr; + float* d_disp_cropped_ = nullptr; + float* d_disp_for_depth_ = nullptr; + int64_t depth_alloc_pixels_ = 0; + + int scaled_w_ = 0; // Uniform-scaled width (before padding) + int scaled_h_ = 0; // Uniform-scaled height (before padding) + + bool gwc_fp16_ = false; // GWC volume tensor is FP16 (vs. FP32) +}; + +} // namespace ffs_depth diff --git a/cv/fast-foundationstereo/cpp/include/ffs_gwc_plugin.hpp b/cv/fast-foundationstereo/cpp/include/ffs_gwc_plugin.hpp new file mode 100644 index 0000000..f0c876a --- /dev/null +++ b/cv/fast-foundationstereo/cpp/include/ffs_gwc_plugin.hpp @@ -0,0 +1,12 @@ +#pragma once + +namespace ffs_depth { + +// Registers the FFSGWCVolume TensorRT plugin creator in the global registry. +// Safe to call more than once. +bool registerFFSGWCPlugin(); + +} // namespace ffs_depth + +// C ABI wrapper for loading/registering the plugin from Python via ctypes. +extern "C" bool ffs_register_gwc_plugin(); diff --git a/cv/fast-foundationstereo/cpp/src/depth_kernels.cu b/cv/fast-foundationstereo/cpp/src/depth_kernels.cu new file mode 100644 index 0000000..822de1a --- /dev/null +++ b/cv/fast-foundationstereo/cpp/src/depth_kernels.cu @@ -0,0 +1,345 @@ +/** + * CUDA kernels for the Fast-FoundationStereo TensorRT depth inference pipeline. + * + * Active kernels: + * 1. GWC volume (mixed precision) -- FP32 features -> FP16 correlation volume + * 2. HWC uint8 -> CHW float -- exact-match path (no resize, with zero-padding) + * 3. Uniform resize + pad -- aspect-ratio-preserving bilinear resize with + * border-replicate padding to model dimensions + * 4. Disparity crop -- removes padding from model-resolution disparity + * 5. Disparity upsample -- nearest-neighbor upsample with scale correction + * 6. Disparity clamp -- clamps to minimum value + * 7. Disparity to depth -- depth_m = fx * baseline_m / disparity (float32, meters) + */ + +#include +#include +#include +#include +#include + +namespace ffs_depth { +namespace cuda { + +// ========================================================================= +// 1. GWC Volume +// ========================================================================= +template +__device__ __forceinline__ float toFloat(T v) { + return static_cast(v); +} + +template <> +__device__ __forceinline__ float toFloat<__half>(__half v) { + return __half2float(v); +} + +template +__device__ __forceinline__ T fromFloat(float v) { + return static_cast(v); +} + +template <> +__device__ __forceinline__ __half fromFloat<__half>(float v) { + return __float2half(v); +} + +template +__global__ void buildGWCVolumeKernel( + const InputT* __restrict__ feat_left, + const InputT* __restrict__ feat_right, + OutputT* __restrict__ gwc_volume, + int B, int C, int H, int W, + int max_disp, int num_groups, bool normalize) +{ + const int w = blockIdx.x * blockDim.x + threadIdx.x; + const int h = blockIdx.y * blockDim.y + threadIdx.y; + const int dgb = blockIdx.z; + + if (w >= W || h >= H) return; + + const int d = dgb % max_disp; + const int g = (dgb / max_disp) % num_groups; + const int b = dgb / (max_disp * num_groups); + if (b >= B) return; + + const int K = C / num_groups; + const int w_right = w - d; + + const int out_idx = ((b * num_groups + g) * max_disp + d) * H * W + h * W + w; + + if (w_right < 0) { gwc_volume[out_idx] = fromFloat(0.0f); return; } + + float dot = 0.f, nl = 0.f, nr = 0.f; + const int left_base = (b * C + g * K) * H * W + h * W + w; + const int right_base = (b * C + g * K) * H * W + h * W + w_right; + const int stride = H * W; + + for (int k = 0; k < K; ++k) { + float l = toFloat(feat_left [left_base + k * stride]); + float r = toFloat(feat_right[right_base + k * stride]); + dot += l * r; nl += l * l; nr += r * r; + } + if (normalize) + gwc_volume[out_idx] = fromFloat(dot / (sqrtf(nl) * sqrtf(nr) + 1e-5f)); + else + gwc_volume[out_idx] = fromFloat(dot); +} + +// ========================================================================= +// 2. Preprocess RGB HWC uint8 -> CHW float (exact-match path, no resize) +// Used when input resolution matches model resolution exactly. +// ========================================================================= +__global__ void preprocessRGBToCHWKernel( + const uint8_t* __restrict__ d_rgb_hwc, + float* __restrict__ d_chw, + int src_h, int src_w, + int dst_h, int dst_w) +{ + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= dst_w || y >= dst_h) return; + + const int hw = dst_h * dst_w; + const int idx = y * dst_w + x; + + if (x < src_w && y < src_h) { + const int si = (y * src_w + x) * 3; + d_chw[idx] = static_cast(d_rgb_hwc[si + 2]); // R + d_chw[hw + idx] = static_cast(d_rgb_hwc[si + 1]); // G + d_chw[2 * hw + idx] = static_cast(d_rgb_hwc[si]); // B + } else { + d_chw[idx] = 0.f; d_chw[hw + idx] = 0.f; d_chw[2 * hw + idx] = 0.f; + } +} + +// ========================================================================= +// 3. Uniform resize + border-replicate padding (aspect-ratio preserving) +// Pixels in [0, scaled_w) x [0, scaled_h) are bilinear-sampled from src. +// Pixels in the padding region replicate the nearest edge pixel. +// ========================================================================= +__global__ void resizeUniformAndPadKernel( + const uint8_t* __restrict__ d_rgb_hwc, + float* __restrict__ d_chw, + int src_h, int src_w, + int scaled_h, int scaled_w, + int dst_h, int dst_w) +{ + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= dst_w || y >= dst_h) return; + + const int hw = dst_h * dst_w; + const int idx = y * dst_w + x; + + const int cx = min(x, scaled_w - 1); + const int cy = min(y, scaled_h - 1); + + const float sx = static_cast(src_w) / scaled_w; + const float sy = static_cast(src_h) / scaled_h; + + const float fx = (cx + 0.5f) * sx - 0.5f; + const float fy = (cy + 0.5f) * sy - 0.5f; + + const int x0 = max(0, min(__float2int_rd(fx), src_w - 1)); + const int y0 = max(0, min(__float2int_rd(fy), src_h - 1)); + const int x1 = min(x0 + 1, src_w - 1); + const int y1 = min(y0 + 1, src_h - 1); + + const float wx = fx - floorf(fx); + const float wy = fy - floorf(fy); + + for (int c = 0; c < 3; ++c) { + const int sc = 2 - c; // RGB channel reorder + float v00 = static_cast(d_rgb_hwc[(y0 * src_w + x0) * 3 + sc]); + float v01 = static_cast(d_rgb_hwc[(y0 * src_w + x1) * 3 + sc]); + float v10 = static_cast(d_rgb_hwc[(y1 * src_w + x0) * 3 + sc]); + float v11 = static_cast(d_rgb_hwc[(y1 * src_w + x1) * 3 + sc]); + float val = (1.f - wy) * ((1.f - wx) * v00 + wx * v01) + + wy * ((1.f - wx) * v10 + wx * v11); + d_chw[c * hw + idx] = val; + } +} + +// ========================================================================= +// 4. Crop disparity (remove border-replicate padding) +// Extracts the valid (scaled_w x scaled_h) region from the +// model-resolution (model_w x model_h) disparity output. +// ========================================================================= +__global__ void cropDisparityKernel( + const float* __restrict__ src, float* __restrict__ dst, + int src_w, int dst_h, int dst_w) +{ + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= dst_w || y >= dst_h) return; + dst[y * dst_w + x] = src[y * src_w + x]; +} + +// ========================================================================= +// 5. Upsample float disparity (nearest neighbor) with scale correction +// Upsamples from cropped model resolution to input resolution. +// disp_scale converts disparity from scaled-pixel units to input-pixel units. +// ========================================================================= +__global__ void upsampleDisparityKernel( + const float* __restrict__ src, float* __restrict__ dst, + int src_w, int src_h, + int dst_w, int dst_h, + float disp_scale) +{ + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= dst_w || y >= dst_h) return; + + const float sx = static_cast(src_w) / dst_w; + const float sy = static_cast(src_h) / dst_h; + + const float fx = (x + 0.5f) * sx - 0.5f; + const float fy = (y + 0.5f) * sy - 0.5f; + + const int x0 = max(0, min(__float2int_rn(fx), src_w - 1)); + const int y0 = max(0, min(__float2int_rn(fy), src_h - 1)); + float val = src[y0 * src_w + x0]; + + dst[y * dst_w + x] = val * disp_scale; +} + +// ========================================================================= +// 6. Clamp disparity to minimum value (removes negative/zero artifacts) +// ========================================================================= +__global__ void clampDisparityKernel( + float* __restrict__ disp, int count, float min_val) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= count) return; + disp[idx] = fmaxf(disp[idx], min_val); +} + +// ========================================================================= +// 7. Disparity to depth: depth_m = fx * baseline_m / disparity (float32 meters) +// +// Matches scripts/run_demo.py behaviour: +// - disparity is assumed already clamped to >= 0 (see clampDisparityKernel) +// - pixels where the right-image correspondence would fall off-image +// (x - disp < 0) are marked invalid by setting disparity to +inf, +// which yields depth = 0 after division. +// - disp == 0 yields depth = +inf (matches numpy's float divide-by-zero). +// - no clip on the upper end; output is float32 meters. +// ========================================================================= +__global__ void dispToDepthKernel( + const float* __restrict__ disp, + float* __restrict__ depth_m, + int height, int width, float fx, float baseline_m) +{ + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= width || y >= height) return; + const int i = y * width + x; + + float d = disp[i]; + if (static_cast(x) - d < 0.0f) { + d = CUDART_INF_F; // remove_invisible: drop pixel via inf -> depth 0 + } + depth_m[i] = fx * baseline_m / d; +} + +// ========================================================================= +// Host wrappers (extern "C" for linkage) +// ========================================================================= +extern "C" { + +void ffsCudaBuildGWCVolumeFloat( + const float* d_fl, const float* d_fr, float* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s) +{ + dim3 blk(16, 16); + dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp); + buildGWCVolumeKernel<<>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize); +} + +void ffsCudaBuildGWCVolumeHalf( + const __half* d_fl, const __half* d_fr, __half* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s) +{ + dim3 blk(16, 16); + dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp); + buildGWCVolumeKernel<__half, __half><<>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize); +} + +void ffsCudaBuildGWCVolumeHalfToFloat( + const __half* d_fl, const __half* d_fr, float* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s) +{ + dim3 blk(16, 16); + dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp); + buildGWCVolumeKernel<__half, float><<>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize); +} + +void ffsCudaBuildGWCVolumeMixed( + const float* d_fl, const float* d_fr, __half* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s) +{ + dim3 blk(16, 16); + dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp); + buildGWCVolumeKernel<<>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize); +} + +void ffsCudaPreprocessRGBToCHW( + const uint8_t* d_rgb, float* d_chw, + int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s) +{ + dim3 blk(32, 16); + dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16); + preprocessRGBToCHWKernel<<>>(d_rgb, d_chw, src_h, src_w, dst_h, dst_w); +} + +void ffsCudaResizeUniformAndPad( + const uint8_t* d_rgb, float* d_chw, + int src_h, int src_w, int scaled_h, int scaled_w, + int dst_h, int dst_w, cudaStream_t s) +{ + dim3 blk(32, 16); + dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16); + resizeUniformAndPadKernel<<>>(d_rgb, d_chw, src_h, src_w, scaled_h, scaled_w, dst_h, dst_w); +} + +void ffsCudaCropDisparity( + const float* d_src, float* d_dst, + int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s) +{ + dim3 blk(32, 16); + dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16); + cropDisparityKernel<<>>(d_src, d_dst, src_w, dst_h, dst_w); +} + +void ffsCudaUpsampleDisparity( + const float* d_src, float* d_dst, + int src_w, int src_h, int dst_w, int dst_h, + float disp_scale, cudaStream_t s) +{ + dim3 blk(32, 16); + dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16); + upsampleDisparityKernel<<>>(d_src, d_dst, src_w, src_h, dst_w, dst_h, disp_scale); +} + +void ffsCudaClampDisparity( + float* d_disp, int count, float min_val, cudaStream_t s) +{ + int threads = 256; + int blocks = (count + threads - 1) / threads; + clampDisparityKernel<<>>(d_disp, count, min_val); +} + +void ffsCudaDispToDepth( + const float* d_disp, float* d_depth_m, + int height, int width, float fx, float baseline_m, cudaStream_t s) +{ + dim3 blk(32, 16); + dim3 grd((width + 31) / 32, (height + 15) / 16); + dispToDepthKernel<<>>(d_disp, d_depth_m, height, width, fx, baseline_m); +} + +} // extern "C" + +} // namespace cuda +} // namespace ffs_depth diff --git a/cv/fast-foundationstereo/cpp/src/ffs_depth_single_tensorrt.cpp b/cv/fast-foundationstereo/cpp/src/ffs_depth_single_tensorrt.cpp new file mode 100644 index 0000000..ebeb482 --- /dev/null +++ b/cv/fast-foundationstereo/cpp/src/ffs_depth_single_tensorrt.cpp @@ -0,0 +1,392 @@ +#include "ffs_depth_single_tensorrt.hpp" + +#include "ffs_gwc_plugin.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ffs_depth { + +namespace cuda { +extern "C" { + +void ffsCudaPreprocessRGBToCHW( + const uint8_t* d_rgb, float* d_chw, + int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s); + +void ffsCudaResizeUniformAndPad( + const uint8_t* d_rgb, float* d_chw, + int src_h, int src_w, int scaled_h, int scaled_w, + int dst_h, int dst_w, cudaStream_t s); + +void ffsCudaCropDisparity( + const float* d_src, float* d_dst, + int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s); + +void ffsCudaUpsampleDisparity( + const float* d_src, float* d_dst, + int src_w, int src_h, int dst_w, int dst_h, + float disp_scale, cudaStream_t s); + +void ffsCudaClampDisparity( + float* d_disp, int count, float min_val, cudaStream_t s); + +void ffsCudaDispToDepth( + const float* d_disp, float* d_depth_m, + int height, int width, float fx, float baseline_m, cudaStream_t s); + +} // extern "C" +} // namespace cuda + +namespace { + +class TrtLogger : public nvinfer1::ILogger { +public: + void log(Severity severity, const char* msg) noexcept override { + if (severity <= Severity::kWARNING) { + std::cerr << "[TRT] " << msg << std::endl; + } + } +}; + +TrtLogger g_trt_logger; + +size_t elementSize(nvinfer1::DataType dt) { + switch (dt) { + case nvinfer1::DataType::kFLOAT: return 4; + case nvinfer1::DataType::kHALF: return 2; + case nvinfer1::DataType::kINT8: return 1; + case nvinfer1::DataType::kINT32: return 4; + default: return 4; + } +} + +void cudaMallocChecked(void** ptr, size_t bytes, const char* what) { + cudaError_t err = cudaMalloc(ptr, bytes); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("[FFS single] cudaMalloc failed for ") + + what + " (" + std::to_string(bytes) + " bytes): " + + cudaGetErrorString(err)); + } +} + +bool hasTensor(nvinfer1::ICudaEngine& engine, const char* name) { + for (int32_t i = 0; i < engine.getNbIOTensors(); ++i) { + if (std::string(engine.getIOTensorName(i)) == name) return true; + } + return false; +} + +std::string findSingleYamlConfig(const std::string& engine_dir) { + namespace fs = std::filesystem; + + std::vector yaml_files; + std::error_code ec; + fs::directory_iterator it(engine_dir, ec); + if (ec) { + throw std::runtime_error("[FFS single] Cannot list engine directory: " + + engine_dir + " (" + ec.message() + ")"); + } + + for (const auto& entry : it) { + if (ec) { + throw std::runtime_error("[FFS single] Cannot list engine directory: " + + engine_dir + " (" + ec.message() + ")"); + } + if (!entry.is_regular_file()) continue; + + std::string ext = entry.path().extension().string(); + std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + if (ext == ".yaml" || ext == ".yml") { + yaml_files.push_back(entry.path()); + } + } + + if (yaml_files.empty()) { + throw std::runtime_error("[FFS single] No YAML config found in: " + engine_dir); + } + if (yaml_files.size() > 1) { + std::ostringstream oss; + oss << "[FFS single] Expected exactly one YAML config in " << engine_dir + << ", found " << yaml_files.size() << ":"; + for (const auto& path : yaml_files) { + oss << " " << path.string(); + } + throw std::runtime_error(oss.str()); + } + + return yaml_files.front().string(); +} + +std::string trim(std::string s) { + const char* ws = " \t\r\n"; + const size_t start = s.find_first_not_of(ws); + if (start == std::string::npos) return ""; + const size_t end = s.find_last_not_of(ws); + return s.substr(start, end - start + 1); +} + +std::vector parseInts(std::string s) { + for (char& c : s) { + if (!(c >= '0' && c <= '9')) c = ' '; + } + std::stringstream ss(s); + std::vector values; + int v = 0; + while (ss >> v) values.push_back(v); + return values; +} + +bool parseBool(std::string s) { + s = trim(s); + std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return s == "true" || s == "1" || s == "yes" || s == "on"; +} + +} // namespace + +FFSSingleEngineInference::FFSSingleEngineInference(const std::string& engine_dir) { + cudaError_t err = cudaStreamCreate(&stream_); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("[FFS single] cudaStreamCreate failed: ") + + cudaGetErrorString(err)); + } + try { + loadConfig(findSingleYamlConfig(engine_dir)); + if (!registerFFSGWCPlugin()) { + throw std::runtime_error("[FFS single] failed to register FFSGWCVolume plugin"); + } + runtime_.reset(nvinfer1::createInferRuntime(g_trt_logger)); + if (!runtime_) { + throw std::runtime_error("[FFS single] createInferRuntime returned null"); + } + loadEngine(engine_dir + "/fast_foundationstereo.engine"); + allocateBuffers(); + } catch (...) { + freeDeviceBuffers(); + if (stream_) { + cudaStreamDestroy(stream_); + stream_ = nullptr; + } + throw; + } +} + +FFSSingleEngineInference::~FFSSingleEngineInference() { + freeDeviceBuffers(); + if (stream_) { + cudaStreamDestroy(stream_); + stream_ = nullptr; + } +} + +void FFSSingleEngineInference::freeDeviceBuffers() { + auto free = [](void*& p) { if (p) { cudaFree(p); p = nullptr; } }; + free(reinterpret_cast(d_left_)); + free(reinterpret_cast(d_right_)); + free(reinterpret_cast(d_disp_)); + free(reinterpret_cast(d_disp_cropped_)); + free(reinterpret_cast(d_disp_for_depth_)); + depth_alloc_pixels_ = 0; +} + +void FFSSingleEngineInference::loadConfig(const std::string& path) { + std::ifstream f(path); + if (!f.good()) throw std::runtime_error("[FFS single] Cannot open config: " + path); + + std::string line; + int image_index = -1; + while (std::getline(f, line)) { + const size_t comment = line.find('#'); + if (comment != std::string::npos) line = line.substr(0, comment); + line = trim(line); + if (line.empty()) continue; + + if (image_index >= 0) { + if (line.rfind("-", 0) == 0) { + const auto values = parseInts(line); + if (!values.empty()) { + if (image_index == 0) config_.image_height = values[0]; + if (image_index == 1) config_.image_width = values[0]; + ++image_index; + if (image_index >= 2) image_index = -1; + continue; + } + } + image_index = -1; + } + + const size_t colon = line.find(':'); + if (colon == std::string::npos) continue; + const std::string key = trim(line.substr(0, colon)); + const std::string value = trim(line.substr(colon + 1)); + + if (key == "image_size") { + const auto values = parseInts(value); + if (values.size() >= 2) { + config_.image_height = values[0]; + config_.image_width = values[1]; + } else { + image_index = 0; + } + } else if (key == "max_disp") { + const auto values = parseInts(value); + if (!values.empty()) config_.max_disp = values[0]; + } else if (key == "cv_group") { + const auto values = parseInts(value); + if (!values.empty()) config_.cv_group = values[0]; + } else if (key == "valid_iters") { + const auto values = parseInts(value); + if (!values.empty()) config_.valid_iters = values[0]; + } else if (key == "normalize") { + config_.normalize = parseBool(value); + } + } +} + +void FFSSingleEngineInference::loadEngine(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f.good()) throw std::runtime_error("[FFS single] Cannot open engine: " + path); + + f.seekg(0, std::ios::end); + size_t sz = f.tellg(); + f.seekg(0, std::ios::beg); + std::vector buf(sz); + f.read(buf.data(), sz); + + engine_.reset(runtime_->deserializeCudaEngine(buf.data(), sz)); + if (!engine_) throw std::runtime_error("[FFS single] Deserialize failed: " + path); + + context_.reset(engine_->createExecutionContext()); + if (!context_) throw std::runtime_error("[FFS single] Context creation failed: " + path); + + if (!hasTensor(*engine_, "left") || !hasTensor(*engine_, "right") || !hasTensor(*engine_, "disp")) { + throw std::runtime_error("[FFS single] Engine must expose tensors named left, right, and disp"); + } +} + +void FFSSingleEngineInference::allocateBuffers() { + const size_t H = static_cast(config_.image_height); + const size_t W = static_cast(config_.image_width); + const size_t input_bytes = 3 * H * W * sizeof(float); + + cudaMallocChecked(reinterpret_cast(&d_left_), input_bytes, "d_left_"); + cudaMallocChecked(reinterpret_cast(&d_right_), input_bytes, "d_right_"); + + const size_t disp_bytes = H * W * elementSize(engine_->getTensorDataType("disp")); + cudaMallocChecked(reinterpret_cast(&d_disp_), disp_bytes, "d_disp_"); + cudaMallocChecked(reinterpret_cast(&d_disp_cropped_), + H * W * sizeof(float), "d_disp_cropped_"); + + context_->setTensorAddress("left", d_left_); + context_->setTensorAddress("right", d_right_); + context_->setTensorAddress("disp", d_disp_); +} + +void FFSSingleEngineInference::preprocessRGBGPU( + const uint8_t* d_rgb, int src_h, int src_w, float* d_output) { + const int mH = config_.image_height; + const int mW = config_.image_width; + if (src_h == mH && src_w == mW) { + scaled_h_ = mH; + scaled_w_ = mW; + cuda::ffsCudaPreprocessRGBToCHW(d_rgb, d_output, src_h, src_w, mH, mW, stream_); + } else { + const float scale = std::min(static_cast(mW) / src_w, + static_cast(mH) / src_h); + scaled_w_ = std::max(1, static_cast(std::round(src_w * scale))); + scaled_h_ = std::max(1, static_cast(std::round(src_h * scale))); + cuda::ffsCudaResizeUniformAndPad( + d_rgb, d_output, src_h, src_w, scaled_h_, scaled_w_, mH, mW, stream_); + } +} + +void FFSSingleEngineInference::infer( + const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float* d_disp_out) { + if (!d_left_rgb || !d_right_rgb || !d_disp_out) { + throw std::runtime_error("[FFS single] infer: null device pointer"); + } + if (input_h <= 0 || input_w <= 0) { + throw std::runtime_error("[FFS single] infer: input dimensions must be positive"); + } + + const int mH = config_.image_height; + const int mW = config_.image_width; + const bool needs_resize = (input_h != mH || input_w != mW); + + preprocessRGBGPU(d_left_rgb, input_h, input_w, d_left_); + preprocessRGBGPU(d_right_rgb, input_h, input_w, d_right_); + + if (!needs_resize) { + context_->setTensorAddress("disp", d_disp_out); + } + if (!context_->enqueueV3(stream_)) { + throw std::runtime_error("[FFS single] enqueue failed"); + } + + if (needs_resize) { + cuda::ffsCudaClampDisparity(d_disp_, mH * mW, 0.0f, stream_); + cuda::ffsCudaCropDisparity(d_disp_, d_disp_cropped_, + mH, mW, scaled_h_, scaled_w_, stream_); + const float disp_scale = static_cast(input_w) / scaled_w_; + cuda::ffsCudaUpsampleDisparity(d_disp_cropped_, d_disp_out, + scaled_w_, scaled_h_, input_w, input_h, + disp_scale, stream_); + } else { + cuda::ffsCudaClampDisparity(d_disp_out, mH * mW, 0.0f, stream_); + context_->setTensorAddress("disp", d_disp_); + } +} + +void FFSSingleEngineInference::dispToDepth( + const float* d_disp, + int height, int width, + float fx, float baseline_m, + float* d_depth_out) { + cuda::ffsCudaDispToDepth(d_disp, d_depth_out, + height, width, fx, baseline_m, stream_); +} + +void FFSSingleEngineInference::inferDepth( + const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float fx, float baseline_m, + float* d_depth_out) { + if (input_h <= 0 || input_w <= 0) { + throw std::runtime_error("[FFS single] inferDepth: input dimensions must be positive"); + } + const size_t num_pixels = static_cast(input_h) * static_cast(input_w); + if (static_cast(num_pixels) > depth_alloc_pixels_) { + if (d_disp_for_depth_) { + cudaFree(d_disp_for_depth_); + d_disp_for_depth_ = nullptr; + } + cudaMallocChecked(reinterpret_cast(&d_disp_for_depth_), + num_pixels * sizeof(float), "d_disp_for_depth_"); + depth_alloc_pixels_ = static_cast(num_pixels); + } + infer(d_left_rgb, d_right_rgb, input_h, input_w, d_disp_for_depth_); + cuda::ffsCudaDispToDepth(d_disp_for_depth_, d_depth_out, + input_h, input_w, fx, baseline_m, stream_); +} + +void FFSSingleEngineInference::sync() { + cudaStreamSynchronize(stream_); +} + +} // namespace ffs_depth diff --git a/cv/fast-foundationstereo/cpp/src/ffs_depth_tensorrt.cpp b/cv/fast-foundationstereo/cpp/src/ffs_depth_tensorrt.cpp new file mode 100644 index 0000000..fb68615 --- /dev/null +++ b/cv/fast-foundationstereo/cpp/src/ffs_depth_tensorrt.cpp @@ -0,0 +1,488 @@ +#include "ffs_depth_tensorrt.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ffs_depth { + +// Forward declarations of CUDA kernel wrappers (defined in depth_kernels.cu) +namespace cuda { +extern "C" { + +void ffsCudaBuildGWCVolumeMixed( + const float* d_fl, const float* d_fr, __half* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s); + +void ffsCudaPreprocessRGBToCHW( + const uint8_t* d_rgb, float* d_chw, + int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s); + +void ffsCudaResizeUniformAndPad( + const uint8_t* d_rgb, float* d_chw, + int src_h, int src_w, int scaled_h, int scaled_w, + int dst_h, int dst_w, cudaStream_t s); + +void ffsCudaCropDisparity( + const float* d_src, float* d_dst, + int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s); + +void ffsCudaUpsampleDisparity( + const float* d_src, float* d_dst, + int src_w, int src_h, int dst_w, int dst_h, + float disp_scale, cudaStream_t s); + +void ffsCudaClampDisparity( + float* d_disp, int count, float min_val, cudaStream_t s); + +void ffsCudaDispToDepth( + const float* d_disp, float* d_depth_m, + int height, int width, float fx, float baseline_m, cudaStream_t s); + +} // extern "C" +} // namespace cuda + +namespace { + +class TrtLogger : public nvinfer1::ILogger { +public: + void log(Severity severity, const char* msg) noexcept override { + if (severity <= Severity::kWARNING) + std::cerr << "[TRT] " << msg << std::endl; + } +}; + +TrtLogger g_trt_logger; + +size_t elementSize(nvinfer1::DataType dt) { + switch (dt) { + case nvinfer1::DataType::kFLOAT: return 4; + case nvinfer1::DataType::kHALF: return 2; + case nvinfer1::DataType::kINT8: return 1; + case nvinfer1::DataType::kINT32: return 4; + default: return 4; + } +} + +void cudaMallocChecked(void** ptr, size_t bytes, const char* what) { + cudaError_t err = cudaMalloc(ptr, bytes); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("[FFS] cudaMalloc failed for ") + what + + " (" + std::to_string(bytes) + " bytes): " + + cudaGetErrorString(err)); + } +} + +std::string trim(std::string s) { + const char* ws = " \t\r\n"; + const size_t start = s.find_first_not_of(ws); + if (start == std::string::npos) return ""; + const size_t end = s.find_last_not_of(ws); + return s.substr(start, end - start + 1); +} + +std::vector parseInts(std::string s) { + for (char& c : s) { + if (!(c >= '0' && c <= '9')) c = ' '; + } + std::stringstream ss(s); + std::vector values; + int v = 0; + while (ss >> v) values.push_back(v); + return values; +} + +bool parseBool(std::string s) { + s = trim(s); + std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return s == "true" || s == "1" || s == "yes" || s == "on"; +} + +} // anonymous namespace + +// ========================================================================= +// Construction / destruction +// ========================================================================= + +FFSDepthInference::FFSDepthInference(const std::string& engine_dir) { + cudaError_t err = cudaStreamCreate(&stream_); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("[FFS] cudaStreamCreate failed: ") + + cudaGetErrorString(err)); + } + try { + loadConfig(engine_dir + "/onnx.yaml"); + runtime_.reset(nvinfer1::createInferRuntime(g_trt_logger)); + if (!runtime_) { + throw std::runtime_error("[FFS] createInferRuntime returned null"); + } + loadEngine(engine_dir + "/feature_runner.engine", feature_engine_, feature_context_); + loadEngine(engine_dir + "/post_runner.engine", post_engine_, post_context_); + allocateBuffers(); + } catch (...) { + // Best-effort cleanup of anything allocated so far. Buffers may be partially + // populated; freeDeviceBuffers() is a no-op for null pointers. + freeDeviceBuffers(); + if (stream_) { + cudaStreamDestroy(stream_); + stream_ = nullptr; + } + throw; + } +} + + +FFSDepthInference::~FFSDepthInference() { + freeDeviceBuffers(); + if (stream_) { + cudaStreamDestroy(stream_); + stream_ = nullptr; + } +} + +void FFSDepthInference::freeDeviceBuffers() { + auto free = [](void*& p) { if (p) { cudaFree(p); p = nullptr; } }; + + free(reinterpret_cast(d_left_)); + free(reinterpret_cast(d_right_)); + free(reinterpret_cast(d_feat_left_04_)); + free(reinterpret_cast(d_feat_left_08_)); + free(reinterpret_cast(d_feat_left_16_)); + free(reinterpret_cast(d_feat_left_32_)); + free(reinterpret_cast(d_feat_right_04_)); + free(reinterpret_cast(d_stem_2x_)); + free(reinterpret_cast(d_gwc_volume_)); + free(reinterpret_cast(d_disp_)); + free(reinterpret_cast(d_disp_cropped_)); + free(reinterpret_cast(d_disp_for_depth_)); + depth_alloc_pixels_ = 0; +} + +// ========================================================================= +// Config / engine loading +// ========================================================================= + +void FFSDepthInference::loadConfig(const std::string& path) { + std::ifstream f(path); + if (!f.good()) throw std::runtime_error("[FFS] Cannot open config: " + path); + + std::string line; + int image_index = -1; + while (std::getline(f, line)) { + const size_t comment = line.find('#'); + if (comment != std::string::npos) line = line.substr(0, comment); + line = trim(line); + if (line.empty()) continue; + + if (image_index >= 0) { + if (line.rfind("-", 0) == 0) { + const auto values = parseInts(line); + if (!values.empty()) { + if (image_index == 0) config_.image_height = values[0]; + if (image_index == 1) config_.image_width = values[0]; + ++image_index; + if (image_index >= 2) image_index = -1; + continue; + } + } + image_index = -1; + } + + const size_t colon = line.find(':'); + if (colon == std::string::npos) continue; + const std::string key = trim(line.substr(0, colon)); + const std::string value = trim(line.substr(colon + 1)); + + if (key == "image_size") { + const auto values = parseInts(value); + if (values.size() >= 2) { + config_.image_height = values[0]; + config_.image_width = values[1]; + } else { + image_index = 0; + } + } else if (key == "max_disp") { + const auto values = parseInts(value); + if (!values.empty()) config_.max_disp = values[0]; + } else if (key == "cv_group") { + const auto values = parseInts(value); + if (!values.empty()) config_.cv_group = values[0]; + } else if (key == "valid_iters") { + const auto values = parseInts(value); + if (!values.empty()) config_.valid_iters = values[0]; + } else if (key == "normalize") { + config_.normalize = parseBool(value); + } + } + + gwc_disp_levels_ = config_.max_disp / 4; +} + +void FFSDepthInference::loadEngine( + const std::string& path, + std::unique_ptr& engine, + std::unique_ptr& ctx) +{ + std::ifstream f(path, std::ios::binary); + if (!f.good()) throw std::runtime_error("[FFS] Cannot open engine: " + path); + + f.seekg(0, std::ios::end); + size_t sz = f.tellg(); + f.seekg(0, std::ios::beg); + std::vector buf(sz); + f.read(buf.data(), sz); + + engine.reset(runtime_->deserializeCudaEngine(buf.data(), sz)); + if (!engine) throw std::runtime_error("[FFS] Deserialize failed: " + path); + + ctx.reset(engine->createExecutionContext()); + if (!ctx) throw std::runtime_error("[FFS] Context creation failed: " + path); +} + +// ========================================================================= +// Buffer allocation +// ========================================================================= + +void FFSDepthInference::allocateBuffers() { + const size_t H = static_cast(config_.image_height); + const size_t W = static_cast(config_.image_width); + const size_t input_bytes = 3 * H * W * sizeof(float); + + cudaMallocChecked(reinterpret_cast(&d_left_), input_bytes, "d_left_"); + cudaMallocChecked(reinterpret_cast(&d_right_), input_bytes, "d_right_"); + + allocateFeatureBuffers(); + allocatePostBuffers(); +} + +void FFSDepthInference::allocateFeatureBuffers() { + auto getDims = [&](const char* name) { + auto d = feature_engine_->getTensorShape(name); + std::vector v(d.nbDims); + for (int i = 0; i < d.nbDims; ++i) v[i] = d.d[i]; + return v; + }; + + auto allocTensor = [&](const char* name, const std::vector& dims) -> float* { + auto dt = feature_engine_->getTensorDataType(name); + size_t bytes = elementSize(dt); + for (int d : dims) { + if (d <= 0) { + throw std::runtime_error(std::string("[FFS] tensor '") + name + + "' has non-positive dimension"); + } + bytes *= static_cast(d); + } + void* p = nullptr; + cudaMallocChecked(&p, bytes, name); + return static_cast(p); + }; + + feat_04_dims_ = getDims("features_left_04"); + feat_08_dims_ = getDims("features_left_08"); + feat_16_dims_ = getDims("features_left_16"); + feat_32_dims_ = getDims("features_left_32"); + stem_2x_dims_ = getDims("stem_2x"); + + d_feat_left_04_ = allocTensor("features_left_04", feat_04_dims_); + d_feat_left_08_ = allocTensor("features_left_08", feat_08_dims_); + d_feat_left_16_ = allocTensor("features_left_16", feat_16_dims_); + d_feat_left_32_ = allocTensor("features_left_32", feat_32_dims_); + d_feat_right_04_ = allocTensor("features_right_04", feat_04_dims_); + d_stem_2x_ = allocTensor("stem_2x", stem_2x_dims_); + + feature_context_->setTensorAddress("left", d_left_); + feature_context_->setTensorAddress("right", d_right_); + feature_context_->setTensorAddress("features_left_04", d_feat_left_04_); + feature_context_->setTensorAddress("features_left_08", d_feat_left_08_); + feature_context_->setTensorAddress("features_left_16", d_feat_left_16_); + feature_context_->setTensorAddress("features_left_32", d_feat_left_32_); + feature_context_->setTensorAddress("features_right_04", d_feat_right_04_); + feature_context_->setTensorAddress("stem_2x", d_stem_2x_); +} + +void FFSDepthInference::allocatePostBuffers() { + const size_t H = static_cast(config_.image_height); + const size_t W = static_cast(config_.image_width); + const size_t H4 = H / 4; + const size_t W4 = W / 4; + + const auto gwc_dt = post_engine_->getTensorDataType("gwc_volume"); + gwc_fp16_ = (gwc_dt == nvinfer1::DataType::kHALF); + const size_t gwc_elem = gwc_fp16_ ? sizeof(__half) : sizeof(float); + const size_t gwc_bytes = + static_cast(config_.cv_group) * + static_cast(gwc_disp_levels_) * H4 * W4 * gwc_elem; + cudaMallocChecked(reinterpret_cast(&d_gwc_volume_), gwc_bytes, "gwc_volume"); + + const auto disp_dt = post_engine_->getTensorDataType("disp"); + const size_t disp_bytes = H * W * elementSize(disp_dt); + cudaMallocChecked(reinterpret_cast(&d_disp_), disp_bytes, "d_disp_"); + + const size_t disp_cropped_bytes = H * W * sizeof(float); + cudaMallocChecked(reinterpret_cast(&d_disp_cropped_), + disp_cropped_bytes, "d_disp_cropped_"); + + post_context_->setTensorAddress("features_left_04", d_feat_left_04_); + post_context_->setTensorAddress("features_left_08", d_feat_left_08_); + post_context_->setTensorAddress("features_left_16", d_feat_left_16_); + post_context_->setTensorAddress("features_left_32", d_feat_left_32_); + post_context_->setTensorAddress("features_right_04", d_feat_right_04_); + post_context_->setTensorAddress("stem_2x", d_stem_2x_); + post_context_->setTensorAddress("gwc_volume", d_gwc_volume_); + post_context_->setTensorAddress("disp", d_disp_); +} + +// ========================================================================= +// Pre-processing +// ========================================================================= + +void FFSDepthInference::preprocessRGBGPU( + const uint8_t* d_rgb, int src_h, int src_w, float* d_output) +{ + const int mH = config_.image_height; + const int mW = config_.image_width; + + if (src_h == mH && src_w == mW) { + scaled_h_ = mH; + scaled_w_ = mW; + cuda::ffsCudaPreprocessRGBToCHW( + d_rgb, d_output, src_h, src_w, mH, mW, stream_); + } else { + const float scale = std::min(static_cast(mW) / src_w, + static_cast(mH) / src_h); + scaled_w_ = std::max(1, static_cast(std::round(src_w * scale))); + scaled_h_ = std::max(1, static_cast(std::round(src_h * scale))); + cuda::ffsCudaResizeUniformAndPad( + d_rgb, d_output, src_h, src_w, scaled_h_, scaled_w_, mH, mW, stream_); + } +} + +// ========================================================================= +// Pipeline stages +// ========================================================================= + +void FFSDepthInference::runFeatureRunner() { + if (!feature_context_->enqueueV3(stream_)) + throw std::runtime_error("[FFS] feature_runner failed"); +} + +void FFSDepthInference::buildGWCVolume() { + int B = feat_04_dims_[0]; + int C = feat_04_dims_[1]; + int H = feat_04_dims_[2]; + int W = feat_04_dims_[3]; + + // Only mixed precision path is used (FP32 features -> FP16 GWC volume) + cuda::ffsCudaBuildGWCVolumeMixed( + d_feat_left_04_, d_feat_right_04_, + reinterpret_cast<__half*>(d_gwc_volume_), + B, C, H, W, gwc_disp_levels_, config_.cv_group, config_.normalize, stream_); + +} + +void FFSDepthInference::runPostRunner() { + if (!post_context_->enqueueV3(stream_)) + throw std::runtime_error("[FFS] post_runner failed"); +} + + +// ========================================================================= +// Public inference entry point +// ========================================================================= + +void FFSDepthInference::infer( + const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float* d_disp_out) +{ + if (!d_left_rgb || !d_right_rgb || !d_disp_out) { + throw std::runtime_error("[FFS] infer: null device pointer"); + } + if (input_h <= 0 || input_w <= 0) { + throw std::runtime_error("[FFS] infer: input dimensions must be positive"); + } + const int mH = config_.image_height; + const int mW = config_.image_width; + const bool needs_resize = (input_h != mH || input_w != mW); + preprocessRGBGPU(d_left_rgb, input_h, input_w, d_left_); + preprocessRGBGPU(d_right_rgb, input_h, input_w, d_right_); + runFeatureRunner(); + buildGWCVolume(); + if (!needs_resize) { + post_context_->setTensorAddress("disp", d_disp_out); + } + runPostRunner(); + if (needs_resize) { + cuda::ffsCudaClampDisparity(d_disp_, mH * mW, 0.0f, stream_); + // Crop padding: model (mW x mH) -> scaled (scaled_w_ x scaled_h_) + cuda::ffsCudaCropDisparity( + d_disp_, d_disp_cropped_, + mH, mW, scaled_h_, scaled_w_, stream_); + // Upsample to input resolution with disparity scale correction + float disp_scale = static_cast(input_w) / scaled_w_; + cuda::ffsCudaUpsampleDisparity( + d_disp_cropped_, d_disp_out, + scaled_w_, scaled_h_, input_w, input_h, + disp_scale, stream_); + } else { + cuda::ffsCudaClampDisparity(d_disp_out, mH * mW, 0.0f, stream_); + post_context_->setTensorAddress("disp", d_disp_); + } +} + +void FFSDepthInference::sync() { + cudaStreamSynchronize(stream_); +} + +void FFSDepthInference::dispToDepth( + const float* d_disp, + int height, int width, + float fx, float baseline_m, + float* d_depth_out) +{ + cuda::ffsCudaDispToDepth(d_disp, d_depth_out, + height, width, fx, baseline_m, stream_); +} + +void FFSDepthInference::inferDepth( + const uint8_t* d_left_rgb, const uint8_t* d_right_rgb, + int input_h, int input_w, + float fx, float baseline_m, + float* d_depth_out) +{ + if (input_h <= 0 || input_w <= 0) { + throw std::runtime_error("[FFS] inferDepth: input dimensions must be positive"); + } + const size_t num_pixels = static_cast(input_h) * static_cast(input_w); + + if (static_cast(num_pixels) > depth_alloc_pixels_) { + // cudaFree is implicitly stream-ordered: it waits for prior work on the + // device before reclaiming the allocation, so this is safe to call here + // even though earlier inferDepth() calls may still be in flight. + if (d_disp_for_depth_) { + cudaFree(d_disp_for_depth_); + d_disp_for_depth_ = nullptr; + } + cudaMallocChecked(reinterpret_cast(&d_disp_for_depth_), + num_pixels * sizeof(float), "d_disp_for_depth_"); + depth_alloc_pixels_ = static_cast(num_pixels); + } + + infer(d_left_rgb, d_right_rgb, input_h, input_w, d_disp_for_depth_); + cuda::ffsCudaDispToDepth(d_disp_for_depth_, d_depth_out, + input_h, input_w, fx, baseline_m, stream_); +} + +} // namespace ffs_depth diff --git a/cv/fast-foundationstereo/cpp/src/gwc_volume_plugin.cpp b/cv/fast-foundationstereo/cpp/src/gwc_volume_plugin.cpp new file mode 100644 index 0000000..414d90b --- /dev/null +++ b/cv/fast-foundationstereo/cpp/src/gwc_volume_plugin.cpp @@ -0,0 +1,325 @@ +#include "ffs_gwc_plugin.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ffs_depth { +namespace cuda { +extern "C" { + +void ffsCudaBuildGWCVolumeFloat( + const float* d_fl, const float* d_fr, float* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s); + +void ffsCudaBuildGWCVolumeHalf( + const __half* d_fl, const __half* d_fr, __half* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s); + +void ffsCudaBuildGWCVolumeHalfToFloat( + const __half* d_fl, const __half* d_fr, float* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s); + +void ffsCudaBuildGWCVolumeMixed( + const float* d_fl, const float* d_fr, __half* d_gwc, + int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s); + +} // extern "C" +} // namespace cuda + +namespace { + +constexpr char kPluginName[] = "FFSGWCVolume"; +constexpr char kPluginVersion[] = "1"; + +struct GWCParams { + int32_t max_disp = 0; // Disparity levels at feature resolution, e.g. max_disp / 4. + int32_t cv_group = 0; + int32_t normalize = 1; +}; + +int32_t fieldToInt(nvinfer1::PluginField const& field, int32_t fallback) { + if (!field.data || field.length <= 0) return fallback; + if (field.type == nvinfer1::PluginFieldType::kINT32) { + return *static_cast(field.data); + } + if (field.type == nvinfer1::PluginFieldType::kINT64) { + return static_cast(*static_cast(field.data)); + } + return fallback; +} + +class FFSGWCVolumePlugin final : public nvinfer1::IPluginV2DynamicExt { +public: + explicit FFSGWCVolumePlugin(GWCParams params) : params_(params) {} + + FFSGWCVolumePlugin(void const* data, size_t length) { + if (data && length == sizeof(GWCParams)) { + std::memcpy(¶ms_, data, sizeof(GWCParams)); + } + } + + nvinfer1::IPluginV2DynamicExt* clone() const noexcept override { + auto* plugin = new FFSGWCVolumePlugin(params_); + plugin->setPluginNamespace(namespace_.c_str()); + return plugin; + } + + char const* getPluginType() const noexcept override { return kPluginName; } + char const* getPluginVersion() const noexcept override { return kPluginVersion; } + int32_t getNbOutputs() const noexcept override { return 1; } + + nvinfer1::DimsExprs getOutputDimensions( + int32_t outputIndex, + nvinfer1::DimsExprs const* inputs, + int32_t nbInputs, + nvinfer1::IExprBuilder& exprBuilder) noexcept override { + nvinfer1::DimsExprs out{}; + if (outputIndex != 0 || nbInputs != 2 || inputs[0].nbDims != 4) { + out.nbDims = -1; + return out; + } + + out.nbDims = 5; + out.d[0] = inputs[0].d[0]; // B + out.d[1] = exprBuilder.constant(params_.cv_group); + out.d[2] = exprBuilder.constant(params_.max_disp); + out.d[3] = inputs[0].d[2]; // H + out.d[4] = inputs[0].d[3]; // W + return out; + } + + bool supportsFormatCombination( + int32_t pos, + nvinfer1::PluginTensorDesc const* inOut, + int32_t nbInputs, + int32_t nbOutputs) noexcept override { + if (nbInputs != 2 || nbOutputs != 1 || pos < 0 || pos >= 3) return false; + if (inOut[pos].format != nvinfer1::TensorFormat::kLINEAR) return false; + + if (pos == 0) { + return inOut[0].type == nvinfer1::DataType::kFLOAT || + inOut[0].type == nvinfer1::DataType::kHALF; + } + if (pos == 1) { + return inOut[1].type == inOut[0].type; + } + return inOut[2].type == nvinfer1::DataType::kFLOAT || + inOut[2].type == nvinfer1::DataType::kHALF; + } + + void configurePlugin( + nvinfer1::DynamicPluginTensorDesc const*, + int32_t, + nvinfer1::DynamicPluginTensorDesc const*, + int32_t) noexcept override {} + + size_t getWorkspaceSize( + nvinfer1::PluginTensorDesc const*, + int32_t, + nvinfer1::PluginTensorDesc const*, + int32_t) const noexcept override { + return 0; + } + + int32_t enqueue( + nvinfer1::PluginTensorDesc const* inputDesc, + nvinfer1::PluginTensorDesc const* outputDesc, + void const* const* inputs, + void* const* outputs, + void*, + cudaStream_t stream) noexcept override { + if (!inputs || !outputs || !inputs[0] || !inputs[1] || !outputs[0]) return 1; + if (inputDesc[0].dims.nbDims != 4) return 1; + + const int B = inputDesc[0].dims.d[0]; + const int C = inputDesc[0].dims.d[1]; + const int H = inputDesc[0].dims.d[2]; + const int W = inputDesc[0].dims.d[3]; + const bool normalize = params_.normalize != 0; + + if (inputDesc[0].type == nvinfer1::DataType::kFLOAT && + outputDesc[0].type == nvinfer1::DataType::kFLOAT) { + cuda::ffsCudaBuildGWCVolumeFloat( + static_cast(inputs[0]), + static_cast(inputs[1]), + static_cast(outputs[0]), + B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream); + return 0; + } + if (inputDesc[0].type == nvinfer1::DataType::kFLOAT && + outputDesc[0].type == nvinfer1::DataType::kHALF) { + cuda::ffsCudaBuildGWCVolumeMixed( + static_cast(inputs[0]), + static_cast(inputs[1]), + static_cast<__half*>(outputs[0]), + B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream); + return 0; + } + if (inputDesc[0].type == nvinfer1::DataType::kHALF && + outputDesc[0].type == nvinfer1::DataType::kHALF) { + cuda::ffsCudaBuildGWCVolumeHalf( + static_cast<__half const*>(inputs[0]), + static_cast<__half const*>(inputs[1]), + static_cast<__half*>(outputs[0]), + B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream); + return 0; + } + if (inputDesc[0].type == nvinfer1::DataType::kHALF && + outputDesc[0].type == nvinfer1::DataType::kFLOAT) { + cuda::ffsCudaBuildGWCVolumeHalfToFloat( + static_cast<__half const*>(inputs[0]), + static_cast<__half const*>(inputs[1]), + static_cast(outputs[0]), + B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream); + return 0; + } + + return 1; + } + + nvinfer1::DataType getOutputDataType( + int32_t, + nvinfer1::DataType const* inputTypes, + int32_t nbInputs) const noexcept override { + if (nbInputs > 0 && inputTypes[0] == nvinfer1::DataType::kHALF) { + return nvinfer1::DataType::kHALF; + } + return nvinfer1::DataType::kFLOAT; + } + + int32_t initialize() noexcept override { return 0; } + void terminate() noexcept override {} + size_t getSerializationSize() const noexcept override { return sizeof(GWCParams); } + void serialize(void* buffer) const noexcept override { + std::memcpy(buffer, ¶ms_, sizeof(GWCParams)); + } + void destroy() noexcept override { delete this; } + void setPluginNamespace(char const* pluginNamespace) noexcept override { + namespace_ = pluginNamespace ? pluginNamespace : ""; + } + char const* getPluginNamespace() const noexcept override { return namespace_.c_str(); } + + void attachToContext(cudnnContext*, cublasContext*, nvinfer1::IGpuAllocator*) noexcept override {} + void detachFromContext() noexcept override {} + +private: + GWCParams params_; + std::string namespace_; +}; + +class FFSGWCVolumePluginCreator final : public nvinfer1::IPluginCreator { +public: + FFSGWCVolumePluginCreator() { + fields_.emplace_back("max_disp", nullptr, nvinfer1::PluginFieldType::kINT32, 1); + fields_.emplace_back("cv_group", nullptr, nvinfer1::PluginFieldType::kINT32, 1); + fields_.emplace_back("normalize", nullptr, nvinfer1::PluginFieldType::kINT32, 1); + field_collection_.nbFields = static_cast(fields_.size()); + field_collection_.fields = fields_.data(); + } + + char const* getPluginName() const noexcept override { return kPluginName; } + char const* getPluginVersion() const noexcept override { return kPluginVersion; } + nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override { + return &field_collection_; + } + + nvinfer1::IPluginV2* createPlugin( + char const*, + nvinfer1::PluginFieldCollection const* fc) noexcept override { + GWCParams params; + if (fc) { + for (int32_t i = 0; i < fc->nbFields; ++i) { + auto const& field = fc->fields[i]; + if (!std::strcmp(field.name, "max_disp")) { + params.max_disp = fieldToInt(field, params.max_disp); + } else if (!std::strcmp(field.name, "cv_group")) { + params.cv_group = fieldToInt(field, params.cv_group); + } else if (!std::strcmp(field.name, "normalize")) { + params.normalize = fieldToInt(field, params.normalize); + } + } + } + if (params.max_disp <= 0 || params.cv_group <= 0) return nullptr; + auto* plugin = new FFSGWCVolumePlugin(params); + plugin->setPluginNamespace(namespace_.c_str()); + return plugin; + } + + nvinfer1::IPluginV2* deserializePlugin( + char const*, + void const* serialData, + size_t serialLength) noexcept override { + auto* plugin = new FFSGWCVolumePlugin(serialData, serialLength); + plugin->setPluginNamespace(namespace_.c_str()); + return plugin; + } + + void setPluginNamespace(char const* pluginNamespace) noexcept override { + namespace_ = pluginNamespace ? pluginNamespace : ""; + } + char const* getPluginNamespace() const noexcept override { + return namespace_.c_str(); + } + +private: + std::vector fields_; + nvinfer1::PluginFieldCollection field_collection_{}; + std::string namespace_; +}; + +FFSGWCVolumePluginCreator g_creator; + +} // namespace + +bool registerFFSGWCPlugin() { + static bool result = false; + static std::once_flag once; + std::call_once(once, []() { + auto try_register = [](nvinfer1::IPluginRegistry* registry, const char* tag) -> bool { + if (!registry) { + std::cerr << "[FFS plugin] " << tag << ": registry unavailable\n"; + return false; + } + if (registry->getPluginCreator(kPluginName, kPluginVersion, "")) { + std::cerr << "[FFS plugin] " << tag + << ": FFSGWCVolume v1 ns=\"\" found-already\n"; + return true; + } + if (registry->registerCreator(g_creator, "")) { + std::cerr << "[FFS plugin] " << tag + << ": FFSGWCVolume v1 ns=\"\" newly-registered\n"; + return true; + } + std::cerr << "[FFS plugin] " << tag + << ": FFSGWCVolume v1 ns=\"\" register FAILED\n"; + return false; + }; + + const bool a = try_register(::getPluginRegistry(), "runtime-registry"); + const bool b = try_register( + nvinfer1::getBuilderPluginRegistry(nvinfer1::EngineCapability::kSTANDARD), + "builder-registry"); + result = a || b; + }); + return result; +} + +// Auto-register the creator with the global TRT plugin registry on shared +// library load. Lets external tools (trtexec --staticPlugins, polygraphy +// --plugins) deserialize engines containing FFSGWCVolume without calling +// registerFFSGWCPlugin() themselves. +REGISTER_TENSORRT_PLUGIN(FFSGWCVolumePluginCreator); + +} // namespace ffs_depth + +extern "C" bool ffs_register_gwc_plugin() { + return ffs_depth::registerFFSGWCPlugin(); +} diff --git a/cv/fast-foundationstereo/demo_data/K.txt b/cv/fast-foundationstereo/demo_data/K.txt new file mode 100644 index 0000000..37902bb --- /dev/null +++ b/cv/fast-foundationstereo/demo_data/K.txt @@ -0,0 +1,2 @@ +754.6680908203125 0.0 489.3794860839844 0.0 754.6680908203125 265.16162109375 0.0 0.0 1.0 +0.063 diff --git a/cv/fast-foundationstereo/demo_data/left.png b/cv/fast-foundationstereo/demo_data/left.png new file mode 100644 index 0000000..0715a0e Binary files /dev/null and b/cv/fast-foundationstereo/demo_data/left.png differ diff --git a/cv/fast-foundationstereo/demo_data/right.png b/cv/fast-foundationstereo/demo_data/right.png new file mode 100644 index 0000000..3874642 Binary files /dev/null and b/cv/fast-foundationstereo/demo_data/right.png differ diff --git a/cv/fast-foundationstereo/docker/dockerfile b/cv/fast-foundationstereo/docker/dockerfile new file mode 100755 index 0000000..25cd557 --- /dev/null +++ b/cv/fast-foundationstereo/docker/dockerfile @@ -0,0 +1,43 @@ +FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 + +ENV TZ=US/Pacific +RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone + +RUN apt-get update --fix-missing && \ + apt-get install -y apt-utils build-essential ca-certificates cmake curl ffmpeg git libturbojpeg-dev pkg-config wget libnvinfer10 libnvinfer-plugin10 libnvonnxparsers10 libnvinfer-dispatch10 libnvinfer-bin zstd libx11-xcb1 libxcb-xinerama0 libxcb-icccm4 libxcb-render-util0 libxcb-shape0 libxcb-keysyms1 libxcb-image0 libxkbcommon-x11-0 libxcb-cursor0 libxcb-xkb1 libxcb-render0 libxcb-shm0 libxcb-sync1 libxcb-xfixes0 libxcb-randr0 libxcb-xtest0 libsm6 libxext6 libxkbcommon0 && \ + rm -rf /var/lib/apt/lists/* + +SHELL ["/bin/bash", "--login", "-c"] + +RUN cd / && wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /miniconda.sh && \ + /bin/bash /miniconda.sh -b -p /opt/conda &&\ + ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh &&\ + echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc &&\ + /bin/bash -c "source ~/.bashrc" && \ + /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main &&\ + /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r &&\ + /opt/conda/bin/conda update -n base -c defaults conda -y &&\ + /opt/conda/bin/conda create -n my python=3.12 + +ENV PATH=$PATH:/opt/conda/envs/my/bin +ENV OPENCV_IO_ENABLE_OPENEXR=1 + +RUN conda init bash &&\ + echo "conda activate my" >> ~/.bashrc &&\ + conda activate my &&\ + pip install uv &&\ + uv pip install torch==2.6.0 torchvision==0.21.0 xformers --index-url https://download.pytorch.org/whl/cu124 + +COPY requirements.txt /tmp/requirements.txt +RUN conda activate my &&\ + uv pip install -r /tmp/requirements.txt &&\ + uv pip install onnxruntime-gpu onnx pycuda cuda-python tensorrt-cu12 tensorrt-lean-cu12 tensorrt-dispatch-cu12 nvidia-modelopt[torch] &&\ + conda install -y -c anaconda h5py &&\ + conda install -y -c conda-forge libstdcxx-ng + +RUN cd / && wget https://developer.download.nvidia.com/compute/tensorrt/10.11.0/local_installers/nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\ + apt install -y ./nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\ + echo 'alias trtexec="/usr/src/tensorrt/bin/trtexec"' >> ~/.bashrc &&\ + rm -f /nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb + +SHELL ["/bin/bash", "-c", "source ~/.bashrc && conda activate my"] diff --git a/cv/fast-foundationstereo/docker/dockerfile_cpp b/cv/fast-foundationstereo/docker/dockerfile_cpp new file mode 100755 index 0000000..3febd80 --- /dev/null +++ b/cv/fast-foundationstereo/docker/dockerfile_cpp @@ -0,0 +1,48 @@ +FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 + +ENV TZ=US/Pacific +RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone + +RUN apt-get update --fix-missing && \ + apt-get install -y apt-utils build-essential ca-certificates cmake curl ffmpeg git libturbojpeg-dev pkg-config wget libnvinfer10 libnvinfer-dev libnvinfer-plugin10 libnvinfer-plugin-dev libnvonnxparsers10 libnvonnxparsers-dev libnvinfer-dispatch10 libnvinfer-bin zstd libyaml-cpp-dev libopencv-dev libx11-xcb1 libxcb-xinerama0 libxcb-icccm4 libxcb-render-util0 libxcb-shape0 libxcb-keysyms1 libxcb-image0 libxkbcommon-x11-0 libxcb-cursor0 libxcb-xkb1 libxcb-render0 libxcb-shm0 libxcb-sync1 libxcb-xfixes0 libxcb-randr0 libxcb-xtest0 libsm6 libxext6 libxkbcommon0 && \ + rm -rf /var/lib/apt/lists/* + +SHELL ["/bin/bash", "--login", "-c"] + +RUN cd / && wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /miniconda.sh && \ + /bin/bash /miniconda.sh -b -p /opt/conda &&\ + ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh &&\ + echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc &&\ + /bin/bash -c "source ~/.bashrc" && \ + /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main &&\ + /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r &&\ + /opt/conda/bin/conda update -n base -c defaults conda -y &&\ + /opt/conda/bin/conda create -n my python=3.12 + +ENV PATH=$PATH:/opt/conda/envs/my/bin +ENV OPENCV_IO_ENABLE_OPENEXR=1 + +RUN conda init bash &&\ + echo "conda activate my" >> ~/.bashrc &&\ + conda activate my &&\ + pip install uv &&\ + uv pip install torch==2.9.1 torchvision==0.24.1 --index-url https://download.pytorch.org/whl/cu128 + +COPY requirements.txt /tmp/requirements.txt +# Do not install nvidia-modelopt[torch] here. It has a strong transitive torch +# requirement and can silently replace the pinned torch/torchvision pair above, +# causing torchvision ABI mismatches such as: +# RuntimeError: operator torchvision::nms does not exist +# Install modelopt manually only if you need quantization tooling. +RUN conda activate my &&\ + uv pip install -r /tmp/requirements.txt &&\ + uv pip install onnxruntime-gpu onnx onnxscript pycuda cuda-python tensorrt-cu12 tensorrt-lean-cu12 tensorrt-dispatch-cu12 && nvidia-modelopt[torch] &&\ + conda install -y -c anaconda h5py &&\ + conda install -y -c conda-forge libstdcxx-ng + +RUN cd / && wget https://developer.download.nvidia.com/compute/tensorrt/10.11.0/local_installers/nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\ + apt install -y ./nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\ + echo 'alias trtexec="/usr/src/tensorrt/bin/trtexec"' >> ~/.bashrc &&\ + rm -f /nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb + +SHELL ["/bin/bash", "-c", "source ~/.bashrc && conda activate my"] diff --git a/cv/fast-foundationstereo/docker/run_container.sh b/cv/fast-foundationstereo/docker/run_container.sh new file mode 100755 index 0000000..a2405c4 --- /dev/null +++ b/cv/fast-foundationstereo/docker/run_container.sh @@ -0,0 +1,4 @@ +docker rm -f ffs || true +xhost +local:root || true +DIR=$(pwd)/../ +docker run --gpus all --runtime nvidia --env NVIDIA_DISABLE_REQUIRE=1 -it --network=host --name ffs --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v $DIR:/workspace --ipc=host -e DISPLAY=${DISPLAY} -v /tmp/.X11-unix:/tmp/.X11-unix -v /tmp:/tmp -v /home:/home -v /mnt:/mnt ffs bash diff --git a/cv/fast-foundationstereo/model_card.md b/cv/fast-foundationstereo/model_card.md new file mode 100755 index 0000000..b1e11a1 --- /dev/null +++ b/cv/fast-foundationstereo/model_card.md @@ -0,0 +1,138 @@ +

Model Card - Fast-FoundationStereo

+ +# Overview + +## Description: +The Fast-FoundationStereo model estimates the disparity of each pixel in a rectified binocular stereo pair of images. This is a transformer based foundational model which shows strong generalization running in real-time. This model is for research and evaluation purposes only. + +### License/Terms of Use: +The code is released using the NVIDIA Source Code License. https://github.com/NVlabs/FoundationStereo/blob/master/LICENSE + +### Deployment Geography: +Global + +### Use Case: +Researchers and developers in the field of computer vision, specifically those interested in depth estimation, are expected to use this method for tasks such as three dimensional reconstruction, object detection, object pose estimation, and scene understanding. + +### Release Date: +Github [02/01/2026] via [https://github.com/NVlabs/Fast-FoundationStereo] + +## Reference(s): +[Fast-FoundationStereo: Real-Time Zero-Shot Stereo Matching](https://arxiv.org/abs/2512.11130) + +## Model Architecture: +**Architecture Type:** Transformers and convolutional neural networks (CNNs). + +**Network Architecture:** The network contains three parts: 1) EdgeNeXt student module that distills the original FoundationStereo feature extractor. 2) Set of blocks (CNNs and transformers) that performs matching with long-range dependencies. 3) Reduced set of convGRU blocks. + +**Number of model parameters:** 14.6M. + + +## Input: +**Input Type(s):** A pair of two rectified binocular stereo images + +**Input Format(s):** Red, Green, Blue (RGB) + +**Input Parameters:** The input parameters to this model are rectified stereo images, specifically Two-Dimensional (2D) images from a camera like Zed. In addition, the baseline is needed to convert disparity to depth. + +**Other Properties Related to Input:** Additional input properties: +- No Alpha Channel or Pre-Processing Needed. Bit: 24-bit. + +## Output: +**Output Type(s):** Disparity image + +**Output Format(s):** 16-bit unsigned integer + +**Output Parameters:** The output parameter of this model is the final 2D disparity map. + +**Other Properties Related to Output:** No Alpha Channel or Post-Processing Needed Bit: 16 bit. + +Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA’s hardware (e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions. + +## Software Integration : +**Runtime Engine(s):** +Not Applicable (N/A) + + +**Supported Hardware Microarchitecture Compatibility:** +NVIDIA Ampere + + +**[Preferred/Supported] Operating System(s):** +Linux + + +## Model Version(s): +v1.0: Initial model version with full capabilities, unpruned and trained. + +# Training and Evaluation Datasets: + +## Training Dataset: + +**Data Modality:** +Image + +**Image Training Data Size:** +1 Million to 1 Billion Images. + + +**Link:** Internal, proprietary dataset, and Stereo4D dataset + +**Data Collection Method by dataset** +[Hybrid: Synthetic, Automatic/Sensors] + +**Labeling Method by dataset** +[Hybrid: Synthetic, Automatic/Sensors] + +**Properties:** The training dataset includes: 1) a large-scale synthetic dataset featuring 1.4 million stereo pairs with large diversity of objects and scenes and high photorealism; 2) real dataset from Stereo4D (external) + +## Testing Dataset: +**Link:** Middlebury dataset + +**Data Collection Method by dataset** +[Automatic/Sensors] + +**Labeling Method by dataset** +[Automatic/Sensors] + + +**Properties:** The dataset encompasses a wide range of scenarios, includes diverse three dimensional assets, captures stereo images under diversely randomized camera parameters, and achieves high fidelity in both rendering and spatial layouts. + +## Evaluation Dataset: +- We evaluated the model using public leaderboards well-known to the stereo community:
+ Middlebury: https://vision.middlebury.edu/stereo/ + + +- This is the classic leaderboard for dense stereo matching, developed at Middlebury College.
+ ETH3D: https://www.eth3d.net/ + + +- This is another popular leaderboard for stereo, developed at ETH.
+ KITTI: https://www.cvlibs.net/datasets/kitti/eval_stereo.php + +**Data Collection Method by dataset** +[Automatic/Sensors] + +** Labeling Method by dataset** +[Automatic/Sensors] + +**Properties:** +* The Middlebury Stereo dataset consists of high-resolution stereo sequences with complex geometry and pixel-accurate ground-truth disparity data. The ground-truth disparities were acquired using a novel technique that employs structured lighting and infrared paint.
+ +* ETHD is a multi-view stereo benchmark / 3D reconstruction benchmark that covers a variety of indoor and outdoor scenes. Ground truth geometry was obtained using a high-precision laser scanner. A DSLR camera as well as a synchronized multi-camera rig with varying field-of-view was used to capture images.
+* KITTI stereo dataset is a cornerstone of autonomous driving research, developed by the Karlsruhe Institute of Technology (KIT) and the Toyota Technological Institute at Chicago (TTIC). It provides real-world, high-resolution stereo imagery paired with precise ground-truth depth data collected from a moving vehicle in diverse urban environments. + + + +## Inference: +**Engine:** Tensor(RT) + +**Test Hardware :** +* Zed Stereo Camera, 3090 + + +## Ethical Considerations: +NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse. + +Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/). + diff --git a/cv/fast-foundationstereo/output_docker/disp_vis.png b/cv/fast-foundationstereo/output_docker/disp_vis.png new file mode 100644 index 0000000..02f0d50 Binary files /dev/null and b/cv/fast-foundationstereo/output_docker/disp_vis.png differ diff --git a/cv/fast-foundationstereo/output_docker/left.png b/cv/fast-foundationstereo/output_docker/left.png new file mode 100644 index 0000000..0715a0e Binary files /dev/null and b/cv/fast-foundationstereo/output_docker/left.png differ diff --git a/cv/fast-foundationstereo/output_docker/right.png b/cv/fast-foundationstereo/output_docker/right.png new file mode 100644 index 0000000..3874642 Binary files /dev/null and b/cv/fast-foundationstereo/output_docker/right.png differ diff --git a/cv/fast-foundationstereo/readme.md b/cv/fast-foundationstereo/readme.md new file mode 100644 index 0000000..990fc03 --- /dev/null +++ b/cv/fast-foundationstereo/readme.md @@ -0,0 +1,216 @@ +# Fast-FoundationStereo: Real-Time Zero-Shot Stereo Matching + +This is the official implementation of our paper accepted to CVPR 2026 + +[[Website]](https://nvlabs.github.io/Fast-FoundationStereo/) [[Paper]](https://arxiv.org/abs/2512.11130) [[Video]](https://www.youtube.com/watch?v=2BUYZojCzXE) + +Authors: [Bowen Wen](https://wenbowen123.github.io/), [Shaurya Dewan](https://www.linkedin.com/in/shaurya-dewan-1b07231a2), [Stan Birchfield](https://research.nvidia.com/person/stan-birchfield) + + +# Abstract +Stereo foundation models achieve strong zero-shot generalization but remain computationally prohibitive for real-time applications. Efficient stereo architectures, on the other hand, sacrifice robustness for speed and require costly per-domain fine-tuning. To bridge this gap, we present Fast-FoundationStereo, a family of architectures that achieve, for the first time, strong zero-shot generalization at real-time frame rate. We employ a divide-and-conquer acceleration strategy with three components: (1) knowledge distillation to compress the hybrid backbone into a single efficient student; (2) blockwise neural architecture search for automatically discovering optimal cost filtering designs under latency budgets, reducing search complexity exponentially; and (3) structured pruning for eliminating redundancy in the iterative refinement module. Furthermore, we introduce an automatic pseudo-labeling pipeline used to curate 1.4M in-the-wild stereo pairs to supplement synthetic training data and facilitate knowledge distillation. The resulting model can run over 10× faster than FoundationStereo while closely matching its zero-shot accuracy, thus establishing a new state-of-the-art among real-time methods. + + [NOTE] This model is designed for real-time applications. For offline computation for the best accuracy, please checkout our earlier work [FoundationStereo](https://github.com/NVlabs/FoundationStereo). + + + +

+ +

+ + + + + + + + + +# Environment setup +- Option 1: Docker +```bash +docker build --network host -t ffs -f docker/dockerfile . +bash docker/run_container.sh +``` + +- Option 2: pip +```bash +conda create -n ffs python=3.12 && conda activate ffs +pip install torch==2.6.0 torchvision==0.21.0 xformers --index-url https://download.pytorch.org/whl/cu124 +pip install -r requirements.txt +``` + + +# Weights and Trade-off +download from [here](https://drive.google.com/drive/folders/1HuTt7UIp7gQsMiDvJwVuWmKpvFzIIMap?usp=drive_link) and put under the folder `weights/` (e.g. `./weights/23-36-37`). Below table compares the differences among some representative models of varying sizes from our trained family. They are sorted from slowest to fastest, with accuracy descending, where runtime is profiled on GPU 3090, image size 640x480. + +To trade-off speed and accuracy, there are two options: +1) Try with different checkpoints. +2) Tune the config flags (see explanations in the "Run demo" section below). + +| Checkpoint | valid_iters | Runtime-Pytorch (ms) | Runtime-TRT (ms) | Peak Memory (MB) | +|---------------|-------------|-------------|-----------------|-----------------| +| `23-36-37` | 8 | 49.4 | 23.4 | 653 | +| `23-36-37` | 4 | 41.1 | 18.4 | 653 | +| `20-26-39` | 8 | 43.6 | 19.4 | 651 | +| `20-26-39` | 4 | 37.5 | 16.4 | 651 | +| `20-30-48` | 8 | 38.4 | 16.6 | 646 | +| `20-30-48` | 4 | 29.3 | 14.0 | 646 | + +# Run demo +``` +python scripts/run_demo.py --model_dir weights/23-36-37/model_best_bp2_serialize.pth --left_file demo_data/left.png --right_file demo_data/right.png --intrinsic_file demo_data/K.txt --out_dir output/ --remove_invisible 0 --denoise_cloud 1 --scale 1 --get_pc 1 --valid_iters 8 --max_disp 192 --zfar 100 +``` +| Flag | Meaning | +|-----------------------------|------------------------------------------------------------------------| +| `--model_dir` | Path to the trained weights/model file | +| `--left_file` | Path to the left image file | +| `--right_file` | Path to the right image file | +| `--intrinsic_file` | Path to the camera intrinsic matrix and baseline file | +| `--out_dir` | Output directory for saving results | +| `--remove_invisible` | Whether to ignore non-overlapping region's depth (0: no, 1: yes) | +| `--denoise_cloud` | Whether to apply denoising to the point cloud (0: no, 1: yes) | +| `--scale` | Image scaling factor | +| `--get_pc` | Obtain point cloud output (0: no, 1: yes) | +| `--valid_iters` | Number of refinement updates during forward pass | +| `--max_disp` | Maximum disparity for volume encoding, 192 should be enough, unless you need to sense very near objects (e.g. <0.1m). Increasing it runs slower and uses more memory. | +| `--zfar` | Maximum depth to include in point cloud | + +Refer to `scripts/run_demo.py` for comprehensive list of flags. + +**Tips:** +- The input left and right images should be rectified and undistorted, which means there should not be fisheye kind of lens distortion and the epipolar lines are horizontal between the left/right images. If you obtain images from stereo cameras such as Zed, they usually have handled this for you. +- Do not swap left and right image. The left image should really be obtained from the left-side camera (objects will appear righter in the image). +- We recommend to use PNG files with no lossy compression +- Our method works best on stereo RGB images. However, we have also tested it on monochrome or IR stereo images (e.g. from RealSense D4XX series) and it works well too. +- To get point cloud for your own data, you need to specify the intrinsics. In the intrinsic file in args, 1st line is the flattened 1x9 intrinsic matrix, 2nd line is the baseline (distance) between the left and right camera, unit in meters. +- The model performs better for image width size <1000. You can run with smaller scale, e.g. `--scale 0.5` to downsize input image, then upsize the output depth to your need with nearest neighbor interpolation. +- For faster inference, you can reduce the input image resolution by e.g. `--scale 0.5`, and reduce refine iterations by e.g. `--valid_iters 4`. +- Note that the 1st time running is slower due to compilation, use a while loop after warm up for live running. + +Expect to see results like below: +- Disparity/Depth: +

+ Disparity Visualization +

+ +- Point cloud: +

+ Point Cloud Visualization +

+ + +# ONNX/TRT +For TRT, we recommend first setup env in docker. + +## Single ONNX + +Export the full model as a single ONNX file. This replaces the Triton GWC kernel with ONNX-compatible ops so no intermediate engine split is needed. + +```bash +python scripts/make_single_onnx.py --model_dir weights/23-36-37/model_best_bp2_serialize.pth --save_path output/ --height 480 --width 640 --valid_iters 8 --max_disp 192 +``` + +| Flag | Meaning | +|-------------------|--------------------------------------------------------------------------| +| `--model_dir` | Path to the trained weights/model file | +| `--save_path` | Directory to save the ONNX model and config | +| `--height` | Input image height, must be divisible by 32. Reduce for faster speed. | +| `--width` | Input image width, must be divisible by 32. Reduce for faster speed. | +| `--valid_iters` | Number of refinement updates during forward pass, reduce for faster speed, but may drop quality | +| `--max_disp` | Maximum disparity for volume encoding, 192 should be enough, unless you need to sense very near objects (e.g. <0.1m). Increasing it runs slower and uses more memory. | +| `--onnx_name` | Base name for the saved ONNX file (default: `fast_foundationstereo`) | + +Then convert to a single TRT engine: +```bash +trtexec --onnx=output/fast_foundationstereo.onnx --saveEngine=output/fast_foundationstereo.engine --fp16 +``` + +To run inference with the single ONNX or TRT engine: +```bash +python scripts/run_demo_single_trt.py --model_dir output/ --left_file demo_data/left.png --right_file demo_data/right.png --intrinsic_file demo_data/K.txt --out_dir output_demo/ --get_pc 1 --remove_invisible 0 --denoise_cloud 1 --zfar 100 +``` + +The script auto-detects `.engine` or `.onnx` files in `--model_dir`. To use a specific file, pass `--model_file` directly. + +| Flag | Meaning | +|-----------------------|--------------------------------------------------------------------------| +| `--model_dir` | Directory containing the .onnx/.engine file and its .yaml config | +| `--model_file` | Explicit path to .onnx or .engine file (overrides auto-search) | +| `--left_file` | Path to the left image file | +| `--right_file` | Path to the right image file | +| `--intrinsic_file` | Path to the camera intrinsic matrix and baseline file | +| `--out_dir` | Output directory for saving results | +| `--remove_invisible` | Whether to ignore non-overlapping region's depth (0: no, 1: yes) | +| `--denoise_cloud` | Whether to apply denoising to the point cloud (0: no, 1: yes) | +| `--get_pc` | Obtain point cloud output (0: no, 1: yes) | +| `--zfar` | Maximum depth (m) to include in point cloud | + +**Note:** The single ONNX model expects **pre-normalized** float32 inputs (ImageNet normalization stripped). The inference script handles this automatically. If integrating into your own pipeline, apply normalization beforehand: +``` +normalized = (pixel - mean) / std +mean = [123.675, 116.28, 103.53] +std = [ 58.395, 57.12, 57.375] +``` + +## Two-stage ONNX + +The original export splits the model into two ONNX files around the Triton GWC kernel, which runs as an intermediate step between the two TRT engines. + +```bash +python scripts/make_onnx.py --model_dir weights/23-36-37/model_best_bp2_serialize.pth --save_path output/ --height 448 --width 640 --valid_iters 8 --max_disp 192 +``` + +| Flag | Meaning | +|-------------------|--------------------------------------------------------------------------| +| `--model_dir` | Path to the trained weights/model file | +| `--save_path` | Directory to save ONNX outputs and zip file | +| `--height` | Input image height, better to be divisible by 32. Reduce image size can increase speed. | +| `--width` | Input image width, better to be divisible by 32. Reduce image size can increase speed. | +| `--valid_iters` | Number of updates during forward pass, reduce it for faster speed, but may drop quality | +| `--max_disp` | Maximum disparity for volume encoding, 192 should be enough, unless you need to sense very near objects (e.g. <0.1m). Increasing it runs slower and uses more memory. | + +Refer to `scripts/make_onnx.py` for a comprehensive list of available flags. + +Then convert from ONNX to TRT: +```bash +trtexec --onnx=output/feature_runner.onnx --saveEngine=output/feature_runner.engine --fp16 --useCudaGraph +trtexec --onnx=output/post_runner.onnx --saveEngine=output/post_runner.engine --fp16 --useCudaGraph +``` + +To use the two-stage TRT for inference: +```bash +python scripts/run_demo_tensorrt.py --onnx_dir output/ --left_file demo_data/left.png --right_file demo_data/right.png --intrinsic_file demo_data/K.txt --out_dir output/ --remove_invisible 0 --denoise_cloud 1 --get_pc 1 --zfar 100 +``` + +# Internet-Scale Pseudo-Labeling +Real-world data offers greater diversity and realism than synthetic data. However, obtaining real stereo images with ground-truth metric depth annotation is notoriously difficult. To address this challenge, we propose an automatic data curation pipeline to generate pseudo-labels on internet-scale stereo images from [Stereo4D](https://stereo4d.github.io/) dataset. **Top:** Pseudo-labeling pipeline on in-the-wild internet stereo data. **Bottom:** Visualization of our generated pseudo-labels. + +

+ +

+ +Below are visualizations of the intermediate results in our pseudo-labeling process. In the rightmost column, green checkmark or red cross denotes whether samples are kept for training or not, based on the percentage of positive pixels in the consistency mask. Our data curation process can automatically discover failures on noisy internet data such as images containing subtitle (bottom), mosaic (2nd last row) and overly challenging samples that are unsuitable for training (top). The final pseudo-labels can also correct erroneous predictions from FoundationStereo on sky regions (5th row). + +

+ +

+ + +The dataset is available at HuggingFace: https://huggingface.co/datasets/nvidia/ffs_stereo4d + +# Citation +```bibtex +@article{wen2026fastfoundationstereo, + title={{Fast-FoundationStereo}: Real-Time Zero-Shot Stereo Matching}, + author={Bowen Wen and Shaurya Dewan and Stan Birchfield}, + journal={CVPR}, + year={2026} +} +``` + +# Contact +Please contact [Bowen Wen](https://wenbowen123.github.io/) (bowenw@nvidia.com) for questions and commercial inquiries. + +# Acknowledgement +We would like to thank Xutong Ren, Karsten Patzwaldt, Yonggan Fu, Saurav Muralidharan, Han Cai, Pavlo Molchanov, Yu Wang, Varun Praveen, Joseph Aribido and Jun Gao for their insightful early discussions for this project. We would also like to thank NVIDIA Isaac and TAO teams for their engineering support and valuable discussions. Thanks to the authors of [FoundationStereo](https://github.com/NVlabs/FoundationStereo), [Selective-IGEV](https://github.com/Windsrain/Selective-Stereo), [Stereo4D](https://github.com/Stereo4d/stereo4d-code) and [RAFT-Stereo](https://github.com/princeton-vl/RAFT-Stereo) for their code release. Finally, thanks to CVPR reviewers and AC for their appreciation of this work and constructive feedback. \ No newline at end of file diff --git a/cv/fast-foundationstereo/requirements.txt b/cv/fast-foundationstereo/requirements.txt new file mode 100644 index 0000000..29feba3 --- /dev/null +++ b/cv/fast-foundationstereo/requirements.txt @@ -0,0 +1,19 @@ +# === Model & inference === +timm +einops +omegaconf +scipy +numpy +scikit-image +opencv-contrib-python +imageio +pyyaml +open3d + +# === TensorRT (optional — uncomment if needed) === +# onnx +# onnxruntime-gpu +# tensorrt-cu12 +# tensorrt-lean-cu12 +# tensorrt-dispatch-cu12 +# nvidia-modelopt[torch] diff --git a/cv/fast-foundationstereo/scripts/build_plugin_trt.py b/cv/fast-foundationstereo/scripts/build_plugin_trt.py new file mode 100644 index 0000000..4cf129f --- /dev/null +++ b/cv/fast-foundationstereo/scripts/build_plugin_trt.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +"""Build a TensorRT engine from the FFSGWCVolume plugin ONNX. + +This mirrors cpp/app/build_single_engine.cpp in Python. The ONNX parser still +needs the custom FFSGWCVolume plugin creator registered before parsing, so the +shared plugin library is auto-detected from cpp/build or can be passed with +--plugin_lib. +""" + +import argparse +import ctypes +import os +from pathlib import Path + + +PLUGIN_NAME = "FFSGWCVolume" +PLUGIN_VERSION = "1" +_LOADED_PLUGIN_LIBS = [] + + +def find_default_plugin_library() -> Path | None: + repo_dir = Path(__file__).resolve().parents[1] + candidates = [ + repo_dir / "cpp" / "build" / "libffs_gwc_plugin.so", + repo_dir / "cpp" / "build" / "lib" / "libffs_gwc_plugin.so", + repo_dir / "cpp" / "build" / "Release" / "libffs_gwc_plugin.so", + ] + for path in candidates: + if path.exists(): + return path + return None + + +def load_plugin_library(path: str) -> None: + """Load an optional shared library that registers FFSGWCVolume.""" + lib = ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL) + _LOADED_PLUGIN_LIBS.append(lib) + + # The current C++ code registers via ffs_depth::registerFFSGWCPlugin(). + # A loadable Python plugin library should expose an extern "C" wrapper with + # one of these names so ctypes can call it without C++ name mangling. + for symbol in ("ffs_register_gwc_plugin", "registerFFSGWCPlugin"): + try: + fn = getattr(lib, symbol) + except AttributeError: + continue + fn.restype = ctypes.c_bool + if not fn(): + raise RuntimeError(f"{symbol}() returned false for {path}") + return + + # Some TensorRT plugin libraries register creators during library load. That + # is not true for this repo's current static C++ helper, but allow it here. + + +def find_plugin_creator(trt) -> bool: + registry = trt.get_plugin_registry() + creator = registry.get_plugin_creator(PLUGIN_NAME, PLUGIN_VERSION, "") + return creator is not None + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Build a TensorRT engine from an ONNX graph containing FFSGWCVolume." + ) + parser.add_argument("plugin_onnx", type=Path, help="Path to plugin ONNX file") + parser.add_argument("output_engine", type=Path, help="Path to write TensorRT engine") + parser.add_argument( + "--plugin_lib", + type=Path, + default=None, + help=( + "Shared library that registers FFSGWCVolume. Defaults to " + "cpp/build/libffs_gwc_plugin.so when present. Pure Python cannot " + "provide this repo's CUDA plugin implementation." + ), + ) + parser.add_argument( + "--fp32", + action="store_true", + help="Disable FP16 builder flag. Default allows FP16 when supported.", + ) + parser.add_argument( + "--workspace-mb", + type=int, + default=4096, + help="TensorRT workspace memory limit in MiB.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if not args.plugin_onnx.exists(): + raise FileNotFoundError(f"ONNX file does not exist: {args.plugin_onnx}") + args.output_engine.parent.mkdir(parents=True, exist_ok=True) + + import tensorrt as trt + + logger = trt.Logger(trt.Logger.INFO) + trt.init_libnvinfer_plugins(logger, "") + + plugin_lib = args.plugin_lib or find_default_plugin_library() + if plugin_lib is not None: + if not plugin_lib.exists(): + raise FileNotFoundError(f"Plugin library does not exist: {plugin_lib}") + load_plugin_library(str(plugin_lib)) + + if not find_plugin_creator(trt): + raise RuntimeError( + f"{PLUGIN_NAME} plugin creator is not registered. " + "Build/load a shared library for cpp/src/gwc_volume_plugin.cpp and pass " + "--plugin_lib, or use cpp/build/ffs_build_single_engine." + ) + + explicit_batch = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + builder = trt.Builder(logger) + network = builder.create_network(explicit_batch) + parser = trt.OnnxParser(network, logger) + + parsed = False + if hasattr(parser, "parse_from_file"): + parsed = parser.parse_from_file(str(args.plugin_onnx)) + else: + parsed = parser.parse(args.plugin_onnx.read_bytes()) + if not parsed: + for i in range(parser.num_errors): + print(parser.get_error(i)) + raise RuntimeError(f"failed to parse ONNX: {args.plugin_onnx}") + + config = builder.create_builder_config() + config.set_memory_pool_limit( + trt.MemoryPoolType.WORKSPACE, int(args.workspace_mb) * 1024 * 1024 + ) + if not args.fp32 and builder.platform_has_fast_fp16: + config.set_flag(trt.BuilderFlag.FP16) + + serialized = builder.build_serialized_network(network, config) + if serialized is None: + raise RuntimeError("build_serialized_network failed") + + args.output_engine.write_bytes(bytes(serialized)) + precision = "FP32" if args.fp32 else "FP16 allowed" + print(f"Built engine: {args.output_engine}") + print(f"Precision: {precision}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cv/fast-foundationstereo/scripts/make_onnx.py b/cv/fast-foundationstereo/scripts/make_onnx.py new file mode 100755 index 0000000..be7e165 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/make_onnx.py @@ -0,0 +1,88 @@ +import warnings, argparse, logging, os, sys,zipfile +os.environ['TORCH_COMPILE_DISABLE'] = '1' +os.environ['TORCHDYNAMO_DISABLE'] = '1' +code_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(f'{code_dir}/../') +import omegaconf, yaml, torch,pdb +from omegaconf import OmegaConf +from core.foundation_stereo import FastFoundationStereo, TrtFeatureRunner, TrtPostRunner, build_gwc_volume_triton +import Utils as U + + +class FoundationStereoOnnx(FastFoundationStereo): + def __init__(self, args): + super().__init__(args) + + @torch.no_grad() + def forward(self, left, right): + """ Removes extra outputs and hyper-parameters """ + with torch.amp.autocast('cuda', enabled=True, dtype=U.AMP_DTYPE): + disp = FastFoundationStereo.forward(self, left, right, iters=self.args.valid_iters, test_mode=True, optimize_build_volume=False) + return disp + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + code_dir = os.path.dirname(os.path.abspath(__file__)) + parser.add_argument('--model_dir', type=str, default=f'{code_dir}/../weights/model_best_bp2_serialize.pth') + parser.add_argument('--save_path', type=str, default=f'/home/bowen/debug/', help='Path to save results.') + parser.add_argument('--height', type=int, default=448) + parser.add_argument('--width', type=int, default=640) + parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass') + parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid") + parser.add_argument('--mixed_precision', default=True, action='store_true', help='use mixed precision') + parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid") + parser.add_argument('--n_downsample', type=int, default=2, help="resolution of the disparity field (1/2^K)") + parser.add_argument('--n_gru_layers', type=int, default=1, help="number of hidden GRU levels") + parser.add_argument('--max_disp', type=int, default=192, help="max disp of geometry encoding volume") + parser.add_argument('--low_memory', type=int, default=1, help='reduce memory usage') + args = parser.parse_args() + os.makedirs(os.path.dirname(args.save_path), exist_ok=True) + + torch.autograd.set_grad_enabled(False) + + model = torch.load(args.model_dir, map_location='cpu', weights_only=False) + model.args.max_disp = args.max_disp + model.args.valid_iters = args.valid_iters + model.cuda().eval() + + feature_runner = TrtFeatureRunner(model) + post_runner = TrtPostRunner(model) + + feature_runner.cuda().eval() + post_runner.cuda().eval() + assert args.height % 32 == 0 and args.width % 32 == 0, "height and width must be divisible by 32" + left_img = torch.randn(1, 3, args.height, args.width).cuda().float()*255 + right_img = torch.randn(1, 3, args.height, args.width).cuda().float()*255 + + torch.onnx.export( + feature_runner, + (left_img, right_img), + args.save_path+'/feature_runner.onnx', + opset_version=17, + input_names = ['left', 'right'], + output_names = ['features_left_04', 'features_left_08', 'features_left_16', 'features_left_32', 'features_right_04', 'stem_2x'], + do_constant_folding=True, + dynamo=False, + ) + + features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x = feature_runner(left_img, right_img) + gwc_volume = build_gwc_volume_triton(features_left_04.half(), features_right_04.half(), args.max_disp//4, model.cv_group) + disp = post_runner(features_left_04.float(), features_left_08.float(), features_left_16.float(), features_left_32.float(), features_right_04.float(), stem_2x.float(), gwc_volume.float()) + + torch.onnx.export( + post_runner, + (features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x, gwc_volume), + args.save_path+'/post_runner.onnx', + opset_version=17, + input_names = ['features_left_04', 'features_left_08', 'features_left_16', 'features_left_32', 'features_right_04', 'stem_2x', 'gwc_volume'], + output_names = ['disp'], + do_constant_folding=True, + dynamo=False, + ) + + with open(f'{args.save_path}/onnx.yaml', 'w') as f: + cfg = OmegaConf.to_container(model.args) + cfg['image_size'] = [args.height, args.width] + yaml.safe_dump(cfg, f) diff --git a/cv/fast-foundationstereo/scripts/make_plugin_onnx.py b/cv/fast-foundationstereo/scripts/make_plugin_onnx.py new file mode 100644 index 0000000..ef61d04 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/make_plugin_onnx.py @@ -0,0 +1,160 @@ +import argparse +import logging +import os +import sys + +os.environ['TORCH_COMPILE_DISABLE'] = '1' +os.environ['TORCHDYNAMO_DISABLE'] = '1' + +code_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(f'{code_dir}/../') + + +def build_parser(): + parser = argparse.ArgumentParser( + description='Export Fast-FoundationStereo as one ONNX with an FFSGWCVolume TensorRT plugin node') + parser.add_argument('--model_dir', type=str, + default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth') + parser.add_argument('--save_path', type=str, default=f'{code_dir}/../output_plugin_onnx') + parser.add_argument('--height', type=int, default=608) + parser.add_argument('--width', type=int, default=960) + parser.add_argument('--valid_iters', type=int, default=8) + parser.add_argument('--max_disp', type=int, default=192) + parser.add_argument('--onnx_name', type=str, default='fast_foundationstereo_plugin.onnx') + return parser + + +if any(arg in ('-h', '--help') for arg in sys.argv[1:]): + build_parser().print_help() + sys.exit(0) + +import torch +import torch.nn as nn +import yaml +from omegaconf import OmegaConf +from torch.onnx import symbolic_helper + +from core.foundation_stereo import TrtFeatureRunner, TrtPostRunner + + +class FFSGWCVolumeOp(torch.autograd.Function): + @staticmethod + def forward(ctx, features_left_04, features_right_04, max_disp, cv_group, normalize): + # ONNX export only needs a tensor with the correct static shape here. + # symbolic() emits the TensorRT plugin node that computes the real volume. + batch, _, height, width = features_left_04.shape + return features_left_04.new_zeros( + (batch, int(cv_group), int(max_disp), height, width) + ) + + @staticmethod + def symbolic(g, features_left_04, features_right_04, max_disp, cv_group, normalize): + def as_int(value): + if isinstance(value, int): + return value + return symbolic_helper._parse_arg(value, 'i') + + max_disp = as_int(max_disp) + cv_group = as_int(cv_group) + normalize = as_int(normalize) + out = g.op( + 'FFSGWCVolume', + features_left_04, + features_right_04, + max_disp_i=int(max_disp), + cv_group_i=int(cv_group), + normalize_i=int(normalize), + ) + sizes = features_left_04.type().sizes() + if sizes is not None and len(sizes) == 4: + out.setType(features_left_04.type().with_sizes( + [sizes[0], int(cv_group), int(max_disp), sizes[2], sizes[3]])) + return out + + +class FastFoundationStereoPluginOnnx(nn.Module): + def __init__(self, model, max_disp_levels, cv_group, normalize): + super().__init__() + self.feature_runner = TrtFeatureRunner(model) + self.post_runner = TrtPostRunner(model) + self.max_disp_levels = int(max_disp_levels) + self.cv_group = int(cv_group) + self.normalize = int(bool(normalize)) + + @torch.no_grad() + def forward(self, left, right): + features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x = ( + self.feature_runner(left, right) + ) + gwc_volume = FFSGWCVolumeOp.apply( + features_left_04, + features_right_04, + self.max_disp_levels, + self.cv_group, + self.normalize, + ) + disp = self.post_runner( + features_left_04.float(), + features_left_08.float(), + features_left_16.float(), + features_left_32.float(), + features_right_04.float(), + stem_2x.float(), + gwc_volume.float(), + ) + return disp + + +if __name__ == '__main__': + args = build_parser().parse_args() + + logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s') + assert args.height % 32 == 0 and args.width % 32 == 0, 'height and width must be divisible by 32' + os.makedirs(args.save_path, exist_ok=True) + torch.autograd.set_grad_enabled(False) + + logging.info('Loading model: %s', args.model_dir) + model = torch.load(args.model_dir, map_location='cpu', weights_only=False) + model.args.max_disp = args.max_disp + model.args.valid_iters = args.valid_iters + model.cuda().eval() + + cv_group = int(getattr(model, 'cv_group', getattr(model.args, 'cv_group', 8))) + normalize = bool(getattr(model.args, 'normalize', True)) + wrapper = FastFoundationStereoPluginOnnx( + model, + max_disp_levels=args.max_disp // 4, + cv_group=cv_group, + normalize=normalize, + ).cuda().eval() + + left = torch.randn(1, 3, args.height, args.width, device='cuda').float() * 255 + right = torch.randn(1, 3, args.height, args.width, device='cuda').float() * 255 + + onnx_name = args.onnx_name if args.onnx_name.endswith('.onnx') else f'{args.onnx_name}.onnx' + onnx_path = os.path.join(args.save_path, onnx_name) + logging.info('Exporting plugin ONNX: %s', onnx_path) + + torch.onnx.export( + wrapper, + (left, right), + onnx_path, + opset_version=17, + input_names=['left', 'right'], + output_names=['disp'], + do_constant_folding=True, + dynamo=False, + ) + + cfg = OmegaConf.to_container(model.args) + cfg['image_size'] = [args.height, args.width] + cfg['cv_group'] = cv_group + cfg['normalize'] = normalize + with open(os.path.join(args.save_path, 'onnx.yaml'), 'w') as f: + yaml.safe_dump(cfg, f) + + logging.info('ONNX model: %s', onnx_path) + logging.info('Config : %s', os.path.join(args.save_path, 'onnx.yaml')) + logging.info('Build with:') + logging.info(' cpp/build/ffs_build_single_engine %s %s', + onnx_path, os.path.join(args.save_path, 'fast_foundationstereo.engine')) diff --git a/cv/fast-foundationstereo/scripts/make_single_onnx.py b/cv/fast-foundationstereo/scripts/make_single_onnx.py new file mode 100644 index 0000000..8591a0c --- /dev/null +++ b/cv/fast-foundationstereo/scripts/make_single_onnx.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Export Fast FoundationStereo as a **single** ONNX model. + +Unlike make_onnx.py (which splits into feature_runner + post_runner with a +Triton GWC kernel in between), this script produces one self-contained ONNX +that can be converted to a single TensorRT engine via trtexec. + +Key design choices: + - The GWC and concat cost volumes are built with ONNX-compatible ops + (pad + slice + stack). The upstream pytorch1 variants use + Tensor.unfold / torch.flip which the ONNX exporter cannot handle. + - ImageNet normalization is STRIPPED from the model so that it can be + applied externally (e.g. via Isaac ROS ImageNormalizeNode). The ONNX + model expects **pre-normalized** float inputs: + pixel = (pixel_0_255 - mean) / std + mean = [123.675, 116.28, 103.53] (ImageNet, in 0-255 scale) + std = [ 58.395, 57.12, 57.375] + - Inputs: left_image (1, 3, H, W) float32, ImageNet-normalised + right_image (1, 3, H, W) float32, ImageNet-normalised + - Output: disparity (1, 1, H, W) float32 + +Usage: + python make_single_onnx.py \\ + --model_dir ../weights/model_best_bp2_serialize.pth \\ + --save_path ./output_single_onnx --height 480 --width 640 + + # Then build a TensorRT engine: + trtexec --onnx=./output_single_onnx/fast_foundationstereo.onnx \\ + --saveEngine=fast_foundationstereo.engine --fp16 +""" + +import argparse +import logging +import os +import sys + +os.environ['TORCH_COMPILE_DISABLE'] = '1' +os.environ['TORCHDYNAMO_DISABLE'] = '1' + +code_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(f'{code_dir}/../') + +import yaml +import torch +import torch.nn as nn +import torch.nn.functional as F +from omegaconf import OmegaConf +import core.foundation_stereo as _fs_module +from core.foundation_stereo import FastFoundationStereo + + +# --------------------------------------------------------------------------- +# ONNX-compatible cost-volume builders +# +# The upstream *_optimized_pytorch1 variants use Tensor.unfold + torch.flip +# which the ONNX tracer cannot export. These replacements build the +# disparity-shifted target volume with an explicit loop over disparities +# using only F.pad, slicing, and torch.stack — all fully ONNX-exportable. +# --------------------------------------------------------------------------- + +def _build_gwc_volume_onnx(refimg_fea, targetimg_fea, maxdisp, + num_groups, normalize=True): + dtype = refimg_fea.dtype + B, C, H, W = refimg_fea.shape + channels_per_group = C // num_groups + + ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W) + + shifted = [ + F.pad(targetimg_fea, (d, 0, 0, 0), 'constant', 0.0)[:, :, :, :W] + for d in range(maxdisp) + ] + target_volume = torch.stack(shifted, dim=2) + + ref_volume = ref_volume.view(B, num_groups, channels_per_group, + maxdisp, H, W) + target_volume = target_volume.view(B, num_groups, channels_per_group, + maxdisp, H, W) + + if normalize: + ref_volume = F.normalize(ref_volume.float(), dim=2).to(dtype) + target_volume = F.normalize(target_volume.float(), dim=2).to(dtype) + + return (ref_volume * target_volume).sum(dim=2).contiguous() + + +def _build_concat_volume_onnx(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + + ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W) + + shifted = [ + F.pad(targetimg_fea, (d, 0, 0, 0), 'constant', 0.0)[:, :, :, :W] + for d in range(maxdisp) + ] + target_volume = torch.stack(shifted, dim=2) + + return torch.cat((ref_volume, target_volume), dim=1).contiguous() + + +# --------------------------------------------------------------------------- + +class FastFoundationStereoSingleOnnx(nn.Module): + """Thin wrapper that calls the full model with ONNX-compatible settings. + + Before ONNX tracing the caller must monkey-patch: + - normalize_image → identity (normalization done externally) + - build_gwc_volume_* → _build_gwc_volume_onnx + - build_concat_volume_* → _build_concat_volume_onnx + """ + + def __init__(self, model: FastFoundationStereo): + super().__init__() + self.model = model + + @torch.no_grad() + def forward(self, left_image, right_image): + return self.model.forward( + left_image, right_image, + iters=self.model.args.valid_iters, + test_mode=True, + optimize_build_volume='pytorch1', + ) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description='Export Fast FoundationStereo as a single ONNX model') + parser.add_argument( + '--model_dir', type=str, + default=f'{code_dir}/../weights/model_best_bp2_serialize.pth', + help='Path to the serialized .pth model') + parser.add_argument( + '--save_path', type=str, + default=f'{code_dir}/output_single_onnx', + help='Directory to save the ONNX model and config') + parser.add_argument('--height', type=int, default=480) + parser.add_argument('--width', type=int, default=640) + parser.add_argument('--valid_iters', type=int, default=8, + help='GRU refinement iterations') + parser.add_argument('--max_disp', type=int, default=192, + help='Maximum disparity (in pixels at full resolution)') + parser.add_argument('--onnx_name', type=str, default='fast_foundationstereo', + help='Base name for the saved ONNX file (without .onnx extension)') + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s %(levelname)s %(message)s') + + assert args.height % 32 == 0 and args.width % 32 == 0, \ + 'height and width must be divisible by 32' + + os.makedirs(args.save_path, exist_ok=True) + torch.autograd.set_grad_enabled(False) + + if not os.path.isfile(args.model_dir): + raise FileNotFoundError(f'Model file not found: {args.model_dir}') + + logging.info(f'Loading model from {args.model_dir}') + model = torch.load(args.model_dir, map_location='cpu', weights_only=False) + model.args.max_disp = args.max_disp + model.args.valid_iters = args.valid_iters + model.args.mixed_precision = False + model.cuda().eval() + + wrapper = FastFoundationStereoSingleOnnx(model) + wrapper.cuda().eval() + + left_img = torch.randn(1, 3, args.height, args.width, device='cuda') + right_img = torch.randn(1, 3, args.height, args.width, device='cuda') + + onnx_name = args.onnx_name if args.onnx_name.endswith('.onnx') else f'{args.onnx_name}.onnx' + onnx_path = os.path.join(args.save_path, onnx_name) + logging.info(f'Exporting ONNX ({args.height}x{args.width}) → {onnx_path}') + + # Monkey-patch non-ONNX-exportable functions before tracing + _fs_module.normalize_image = lambda img: img + _fs_module.build_gwc_volume_optimized_pytorch1 = _build_gwc_volume_onnx + _fs_module.build_concat_volume_optimized_pytorch1 = _build_concat_volume_onnx + + torch.onnx.export( + wrapper, + (left_img, right_img), + onnx_path, + opset_version=17, + input_names=['left_image', 'right_image'], + output_names=['disparity'], + do_constant_folding=True, + ) + + cfg = OmegaConf.to_container(model.args) + cfg['image_size'] = [args.height, args.width] + config_name = os.path.splitext(onnx_name)[0] + '.yaml' + config_path = os.path.join(args.save_path, config_name) + with open(config_path, 'w') as f: + yaml.safe_dump(cfg, f) + + logging.info(f'ONNX model : {onnx_path}') + logging.info(f'Config : {config_path}') + logging.info(f'Resolution : {args.height} x {args.width}') + logging.info( + f'Build TRT engine:\n' + f' trtexec --onnx={onnx_path} ' + f'--saveEngine={args.save_path}/{os.path.splitext(onnx_name)[0]}.engine --fp16') diff --git a/cv/fast-foundationstereo/scripts/profile_memory.py b/cv/fast-foundationstereo/scripts/profile_memory.py new file mode 100644 index 0000000..2bb3147 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/profile_memory.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from omegaconf import OmegaConf +from core.utils.utils import InputPadder +import argparse, torch, logging, yaml, time +import numpy as np +from Utils import AMP_DTYPE, set_logging_format, set_seed + + +if __name__=="__main__": + code_dir = os.path.dirname(os.path.realpath(__file__)) + parser = argparse.ArgumentParser() + parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str) + parser.add_argument('--hiera', default=0, type=int) + parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass') + parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity') + parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations') + parser.add_argument('--total', type=int, default=30, help='total number of iterations') + args = parser.parse_args() + + set_logging_format() + set_seed(0) + torch.backends.cudnn.benchmark = True + torch.autograd.set_grad_enabled(False) + + with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff: + cfg:dict = yaml.safe_load(ff) + for k in args.__dict__: + if args.__dict__[k] is not None: + cfg[k] = args.__dict__[k] + args = OmegaConf.create(cfg) + model = torch.load(args.model_dir, map_location='cpu', weights_only=False) + model.args.valid_iters = args.valid_iters + model.args.max_disp = args.max_disp + model.cuda().eval() + + H, W = 480, 640 + img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda() + img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda() + padder = InputPadder(img0.shape, divis_by=32, force_square=False) + img0, img1 = padder.pad(img0, img1) + + logging.info(f"Image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}") + + times = [] + peak_memories = [] + with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE): + for i in range(args.total): + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + t0 = time.perf_counter() + disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='triton') + torch.cuda.synchronize() + elapsed = time.perf_counter() - t0 + peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 2) + times.append(elapsed) + peak_memories.append(peak_mem) + logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms, peak mem: {peak_mem:.1f} MB {'(warmup)' if i < args.warmup else ''}") + + measure_times = times[args.warmup:] + measure_mems = peak_memories[args.warmup:] + avg = np.mean(measure_times) * 1000 + avg_mem = np.mean(measure_mems) + max_mem = np.max(measure_mems) + logging.info(f"Vanilla Pytorch speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters") + logging.info(f"Peak GPU memory (after warmup): avg {avg_mem:.1f} MB, max {max_mem:.1f} MB") diff --git a/cv/fast-foundationstereo/scripts/profile_speed.py b/cv/fast-foundationstereo/scripts/profile_speed.py new file mode 100644 index 0000000..406d00d --- /dev/null +++ b/cv/fast-foundationstereo/scripts/profile_speed.py @@ -0,0 +1,67 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from omegaconf import OmegaConf +from core.utils.utils import InputPadder +import argparse, torch, logging, yaml, time +import numpy as np +from Utils import AMP_DTYPE, set_logging_format, set_seed + + +if __name__=="__main__": + code_dir = os.path.dirname(os.path.realpath(__file__)) + parser = argparse.ArgumentParser() + parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str) + parser.add_argument('--hiera', default=0, type=int) + parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass') + parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity') + parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations') + parser.add_argument('--total', type=int, default=30, help='total number of iterations') + args = parser.parse_args() + + set_logging_format() + set_seed(0) + torch.backends.cudnn.benchmark = True + torch.autograd.set_grad_enabled(False) + + with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff: + cfg:dict = yaml.safe_load(ff) + for k in args.__dict__: + if args.__dict__[k] is not None: + cfg[k] = args.__dict__[k] + args = OmegaConf.create(cfg) + model = torch.load(args.model_dir, map_location='cpu', weights_only=False) + model.args.valid_iters = args.valid_iters + model.args.max_disp = args.max_disp + model.cuda().eval() + + H, W = 480, 640 + img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda() + img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda() + padder = InputPadder(img0.shape, divis_by=32, force_square=False) + img0, img1 = padder.pad(img0, img1) + + logging.info(f"Image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}") + + times = [] + with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE): + for i in range(args.total): + torch.cuda.synchronize() + t0 = time.perf_counter() + disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='triton') + torch.cuda.synchronize() + elapsed = time.perf_counter() - t0 + times.append(elapsed) + logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms {'(warmup)' if i < args.warmup else ''}") + + measure_times = times[args.warmup:] + avg = np.mean(measure_times) * 1000 + logging.info(f"Vanilla Pytorch speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters") diff --git a/cv/fast-foundationstereo/scripts/profile_speed_tensorrt.py b/cv/fast-foundationstereo/scripts/profile_speed_tensorrt.py new file mode 100644 index 0000000..b0895a1 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/profile_speed_tensorrt.py @@ -0,0 +1,66 @@ +import os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from omegaconf import OmegaConf +import argparse, torch, logging, yaml, time +import numpy as np +from Utils import set_logging_format, set_seed +from core.foundation_stereo import TrtRunner + + +def resolve_onnx_cfg_path(onnx_dir: str) -> str: + direct = os.path.join(onnx_dir, 'onnx.yaml') + if os.path.exists(direct): + return direct + parent = os.path.join(os.path.dirname(onnx_dir), 'onnx.yaml') + if os.path.exists(parent): + return parent + raise FileNotFoundError(f"onnx.yaml not found in {onnx_dir} or its parent directory") + + +if __name__=="__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--model_dir', default=None, type=str) + parser.add_argument('--hiera', default=None, type=int) + parser.add_argument('--valid_iters', type=int, default=None, help='number of flow-field updates during forward pass (default: from onnx.yaml)') + parser.add_argument('--max_disp', type=int, default=None, help='maximum disparity (default: from onnx.yaml)') + parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations') + parser.add_argument('--total', type=int, default=30, help='total number of iterations') + parser.add_argument('--build_volume_backend', default=None, choices=['pytorch1', 'triton'], help='backend for cost-volume build (default: from onnx.yaml)') + parser.add_argument('--onnx_dir', default=f'{code_dir}/../output', type=str, help='directory containing TensorRT engines and onnx.yaml') + args = parser.parse_args() + + set_logging_format() + set_seed(0) + torch.backends.cudnn.benchmark = True + torch.autograd.set_grad_enabled(False) + + cfg_path = resolve_onnx_cfg_path(args.onnx_dir) + with open(cfg_path, 'r') as ff: + cfg:dict = yaml.safe_load(ff) + for k in args.__dict__: + if args.__dict__[k] is not None: + cfg[k] = args.__dict__[k] + args = OmegaConf.create(cfg) + + model = TrtRunner(args, args.onnx_dir+'/feature_runner.engine', args.onnx_dir+'/post_runner.engine') + + H, W = int(args.image_size[0]), int(args.image_size[1]) + img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda() + img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda() + + logging.info(f"TensorRT image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}") + + times = [] + for i in range(args.total): + torch.cuda.synchronize() + t0 = time.perf_counter() + _ = model.forward(img0, img1) + torch.cuda.synchronize() + elapsed = time.perf_counter() - t0 + times.append(elapsed) + logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms {'(warmup)' if i < args.warmup else ''}") + + measure_times = times[args.warmup:] + avg = np.mean(measure_times) * 1000 + logging.info(f"TensorRT speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters") diff --git a/cv/fast-foundationstereo/scripts/run_demo.py b/cv/fast-foundationstereo/scripts/run_demo.py new file mode 100755 index 0000000..ef75362 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/run_demo.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from omegaconf import OmegaConf +from core.utils.utils import InputPadder +import argparse, torch, imageio, logging, yaml +import numpy as np +from Utils import ( + AMP_DTYPE, set_logging_format, set_seed, vis_disparity, + depth2xyzmap, toOpen3dCloud, o3d, +) +import cv2 + + +if __name__=="__main__": + code_dir = os.path.dirname(os.path.realpath(__file__)) + parser = argparse.ArgumentParser() + parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str) + parser.add_argument('--left_file', default=f'{code_dir}/../demo_data/left.png', type=str) + parser.add_argument('--right_file', default=f'{code_dir}/../demo_data/right.png', type=str) + parser.add_argument('--intrinsic_file', default=f'{code_dir}/../demo_data/K.txt', type=str, help='camera intrinsic matrix and baseline file') + parser.add_argument('--out_dir', default='/home/bowen/debug/stereo_output', type=str) + parser.add_argument('--remove_invisible', default=1, type=int) + parser.add_argument('--denoise_cloud', default=0, type=int) + parser.add_argument('--denoise_nb_points', type=int, default=30, help='number of points to consider for radius outlier removal') + parser.add_argument('--denoise_radius', type=float, default=0.03, help='radius to use for outlier removal') + parser.add_argument('--scale', default=1, type=float) + parser.add_argument('--hiera', default=0, type=int) + parser.add_argument('--get_pc', type=int, default=1, help='save point cloud output') + parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass') + parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity') + parser.add_argument('--zfar', type=float, default=100, help="max depth to include in point cloud") + args = parser.parse_args() + + set_logging_format() + set_seed(0) + torch.autograd.set_grad_enabled(False) + + os.system(f'rm -rf {args.out_dir} && mkdir -p {args.out_dir}') + + with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff: + cfg:dict = yaml.safe_load(ff) + for k in args.__dict__: + if args.__dict__[k] is not None: + cfg[k] = args.__dict__[k] + args = OmegaConf.create(cfg) + logging.info(f"args:\n{args}") + model = torch.load(args.model_dir, map_location='cpu', weights_only=False) + model.args.valid_iters = args.valid_iters + model.args.max_disp = args.max_disp + + model.cuda().eval() + + scale = args.scale + + img0 = imageio.imread(args.left_file) + img1 = imageio.imread(args.right_file) + if len(img0.shape)==2: + img0 = np.tile(img0[...,None], (1,1,3)) + img1 = np.tile(img1[...,None], (1,1,3)) + img0 = img0[...,:3] + img1 = img1[...,:3] + H,W = img0.shape[:2] + + img0 = cv2.resize(img0, fx=scale, fy=scale, dsize=None) + img1 = cv2.resize(img1, dsize=(img0.shape[1], img0.shape[0])) + H,W = img0.shape[:2] + img0_ori = img0.copy() + img1_ori = img1.copy() + logging.info(f"img0: {img0.shape}") + imageio.imwrite(f'{args.out_dir}/left.png', img0) + imageio.imwrite(f'{args.out_dir}/right.png', img1) + + img0 = torch.as_tensor(img0).cuda().float()[None].permute(0,3,1,2) + img1 = torch.as_tensor(img1).cuda().float()[None].permute(0,3,1,2) + padder = InputPadder(img0.shape, divis_by=32, force_square=False) + img0, img1 = padder.pad(img0, img1) + + logging.info(f"Start forward, 1st time run can be slow due to compilation") + with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE): + if not args.hiera: + disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='pytorch1') + else: + disp = model.run_hierachical(img0, img1, iters=args.valid_iters, test_mode=True, small_ratio=0.5) + logging.info("forward done") + disp = padder.unpad(disp.float()) + disp = disp.data.cpu().numpy().reshape(H,W).clip(0, None) + + cmap = None + min_val = None + max_val = None + vis = vis_disparity(disp, min_val=min_val, max_val=max_val, cmap=cmap, color_map=cv2.COLORMAP_TURBO) + vis = np.concatenate([img0_ori, img1_ori, vis], axis=1) + imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis) + s = 1280/vis.shape[1] + resized_vis = cv2.resize(vis, (int(vis.shape[1]*s), int(vis.shape[0]*s))) + cv2.imshow('disp', resized_vis[:,:,::-1]) + cv2.waitKey(0) + + if args.remove_invisible: + yy,xx = np.meshgrid(np.arange(disp.shape[0]), np.arange(disp.shape[1]), indexing='ij') + us_right = xx-disp + invalid = us_right<0 + disp[invalid] = np.inf + + if args.get_pc: + with open(args.intrinsic_file, 'r') as f: + lines = f.readlines() + K = np.array(list(map(float, lines[0].rstrip().split()))).astype(np.float32).reshape(3,3) + baseline = float(lines[1]) + K[:2] *= scale + depth = K[0,0]*baseline/disp + np.save(f'{args.out_dir}/depth_meter.npy', depth) + xyz_map = depth2xyzmap(depth, K) + pcd = toOpen3dCloud(xyz_map.reshape(-1,3), img0_ori.reshape(-1,3)) + keep_mask = (np.asarray(pcd.points)[:,2]>0) & (np.asarray(pcd.points)[:,2]<=args.zfar) + keep_ids = np.arange(len(np.asarray(pcd.points)))[keep_mask] + pcd = pcd.select_by_index(keep_ids) + o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd) + logging.info(f"PCL saved to {args.out_dir}") + + if args.denoise_cloud: + logging.info("[Optional step] denoise point cloud...") + pcd = pcd.voxel_down_sample(voxel_size=0.001) + cl, ind = pcd.remove_radius_outlier(nb_points=args.denoise_nb_points, radius=args.denoise_radius) + inlier_cloud = pcd.select_by_index(ind) + o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', inlier_cloud) + pcd = inlier_cloud + + logging.info("Visualizing point cloud. Press ESC to exit.") + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.add_geometry(pcd) + vis.get_render_option().point_size = 1.0 + vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5]) + ctr = vis.get_view_control() + ctr.set_front([0, 0, -1]) + id = np.asarray(pcd.points)[:,2].argmin() + ctr.set_lookat(np.asarray(pcd.points)[id]) + ctr.set_up([0, -1, 0]) + vis.run() + vis.destroy_window() diff --git a/cv/fast-foundationstereo/scripts/run_demo_plugin_trt.py b/cv/fast-foundationstereo/scripts/run_demo_plugin_trt.py new file mode 100644 index 0000000..f257360 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/run_demo_plugin_trt.py @@ -0,0 +1,380 @@ +#!/usr/bin/env python3 +"""Run the FFSGWCVolume-plugin TensorRT engine from Python. + +This is the Python equivalent of cpp/src/ffs_single_depth_inference.cpp plus +cpp/app/main.cpp for the single-engine plugin path: + + - engine inputs: left, right + - engine output: disp + - input tensors are RGB, CHW, float32, raw 0-255 values + - resize is aspect-ratio-preserving with right/bottom replicate padding + - disparity is cropped, nearest-neighbor upsampled, and scaled back to the + original input-image pixel units +""" + +import argparse +import logging +import os +import sys +from pathlib import Path + +import cv2 +import imageio.v2 as imageio +import numpy as np +import torch +import yaml + +code_dir = Path(__file__).resolve().parent +repo_dir = code_dir.parent +sys.path.append(str(repo_dir)) + +from Utils import set_logging_format, set_seed, vis_disparity +from build_plugin_trt import ( + PLUGIN_NAME, + find_default_plugin_library, + find_plugin_creator, + load_plugin_library, +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Run Fast-FoundationStereo single TensorRT plugin engine." + ) + parser.add_argument( + "--model_dir", + type=Path, + default=repo_dir / "engine1_plug", + help="Directory containing fast_foundationstereo.engine and one YAML config.", + ) + parser.add_argument( + "--engine_file", + "--model_file", + dest="engine_file", + type=Path, + default=None, + help="Explicit TensorRT engine path. Defaults to /fast_foundationstereo.engine.", + ) + parser.add_argument( + "--config_file", + type=Path, + default=None, + help="Explicit YAML config path. Defaults to the single YAML in the engine directory.", + ) + parser.add_argument( + "--plugin_lib", + type=Path, + default=None, + help="Path to libffs_gwc_plugin.so. Defaults to cpp/build/libffs_gwc_plugin.so.", + ) + parser.add_argument("--left_file", type=Path, default=repo_dir / "demo_data" / "left.png") + parser.add_argument("--right_file", type=Path, default=repo_dir / "demo_data" / "right.png") + parser.add_argument( + "--intrinsic_file", + type=Path, + default=repo_dir / "demo_data" / "K.txt", + help="Text file with 3x3 K on line 1 and baseline in meters on line 2.", + ) + parser.add_argument("--out_dir", type=Path, default=repo_dir / "output_plugin_trt") + return parser.parse_args() + + +def resolve_engine_path(model_dir: Path, engine_file: Path | None) -> Path: + path = engine_file if engine_file is not None else model_dir / "fast_foundationstereo.engine" + if not path.exists(): + raise FileNotFoundError(f"TensorRT engine does not exist: {path}") + return path + + +def resolve_config_path(engine_path: Path, model_dir: Path, config_file: Path | None) -> Path: + if config_file is not None: + if not config_file.exists(): + raise FileNotFoundError(f"Config file does not exist: {config_file}") + return config_file + + search_dir = model_dir if model_dir.exists() else engine_path.parent + yaml_files = sorted( + p for p in search_dir.iterdir() + if p.is_file() and p.suffix.lower() in (".yaml", ".yml") + ) + if not yaml_files: + raise FileNotFoundError(f"No YAML config found in: {search_dir}") + if len(yaml_files) > 1: + names = " ".join(str(p) for p in yaml_files) + raise RuntimeError( + f"Expected exactly one YAML config in {search_dir}, found {len(yaml_files)}: {names}" + ) + return yaml_files[0] + + +def load_intrinsics(path: Path) -> tuple[np.ndarray, float]: + with path.open("r") as f: + lines = f.readlines() + if len(lines) < 2: + raise RuntimeError("intrinsic file must contain K on line 1 and baseline on line 2") + k = np.array(list(map(float, lines[0].strip().split())), dtype=np.float32).reshape(3, 3) + baseline = float(lines[1]) + if not (k[0, 0] > 0.0 and baseline > 0.0): + raise RuntimeError("invalid focal length or baseline in intrinsic file") + return k, baseline + + +def load_rgb(path: Path) -> np.ndarray: + img = imageio.imread(path) + if img.ndim == 2: + img = np.tile(img[..., None], (1, 1, 3)) + img = img[..., :3] + if img.dtype != np.uint8: + img = np.clip(img, 0, 255).astype(np.uint8) + return np.ascontiguousarray(img) + + +def resize_uniform_and_pad_rgb(img: np.ndarray, target_h: int, target_w: int): + src_h, src_w = img.shape[:2] + if src_h == target_h and src_w == target_w: + return img.copy(), target_h, target_w + + scale = min(float(target_w) / src_w, float(target_h) / src_h) + scaled_w = max(1, int(round(src_w * scale))) + scaled_h = max(1, int(round(src_h * scale))) + + scaled = cv2.resize(img, (scaled_w, scaled_h), interpolation=cv2.INTER_LINEAR) + padded = cv2.copyMakeBorder( + scaled, + 0, + target_h - scaled_h, + 0, + target_w - scaled_w, + cv2.BORDER_REPLICATE, + ) + return np.ascontiguousarray(padded), scaled_h, scaled_w + + +def tensor_from_rgb_255(img: np.ndarray) -> torch.Tensor: + arr = img.astype(np.float32, copy=False) + return torch.as_tensor(arr, device="cuda").permute(2, 0, 1).unsqueeze(0).contiguous() + + +def cpp_nearest_resize(src: np.ndarray, dst_h: int, dst_w: int) -> np.ndarray: + src_h, src_w = src.shape + xs = np.rint((np.arange(dst_w, dtype=np.float32) + 0.5) * src_w / dst_w - 0.5) + ys = np.rint((np.arange(dst_h, dtype=np.float32) + 0.5) * src_h / dst_h - 0.5) + xs = np.clip(xs.astype(np.int64), 0, src_w - 1) + ys = np.clip(ys.astype(np.int64), 0, src_h - 1) + return src[ys[:, None], xs[None, :]] + + +def postprocess_disparity( + raw_disp: np.ndarray, + input_h: int, + input_w: int, + model_h: int, + model_w: int, + scaled_h: int, + scaled_w: int, +) -> np.ndarray: + disp = raw_disp.reshape(model_h, model_w).astype(np.float32, copy=False) + disp = np.maximum(disp, 0.0) + + if input_h == model_h and input_w == model_w: + return np.ascontiguousarray(disp) + + cropped = disp[:scaled_h, :scaled_w] + upsampled = cpp_nearest_resize(cropped, input_h, input_w) + return np.ascontiguousarray(upsampled * (float(input_w) / scaled_w), dtype=np.float32) + + +def disparity_to_depth_cpp(disp: np.ndarray, fx: float, baseline_m: float) -> np.ndarray: + d = disp.astype(np.float32, copy=True) + xs = np.arange(d.shape[1], dtype=np.float32)[None, :] + d[(xs - d) < 0.0] = np.inf + with np.errstate(divide="ignore", invalid="ignore"): + depth = np.float32(fx * baseline_m) / d + return np.ascontiguousarray(depth, dtype=np.float32) + + +def save_float_matrix_bin(path: Path, values: np.ndarray) -> None: + arr = np.ascontiguousarray(values, dtype=np.float32) + dims = np.array(arr.shape[:2], dtype=np.int32) + with path.open("wb") as f: + dims.tofile(f) + arr.tofile(f) + + +def colorize_depth(depth: np.ndarray) -> np.ndarray: + valid = np.isfinite(depth) & (depth > 0.0) + if not np.any(valid): + return np.zeros((*depth.shape, 3), dtype=np.uint8) + + safe = np.zeros_like(depth, dtype=np.float32) + safe[valid] = depth[valid] + min_val = float(safe[valid].min()) + max_val = float(safe[valid].max()) + if max_val <= min_val: + max_val = min_val + 1.0 + + depth_u8 = np.clip((safe - min_val) * (255.0 / (max_val - min_val)), 0, 255).astype(np.uint8) + colored_bgr = cv2.applyColorMap(depth_u8, cv2.COLORMAP_TURBO) + colored = colored_bgr[..., ::-1] + colored[~valid] = 0 + return colored + + +class PluginTensorRTRunner: + def __init__(self, engine_path: Path, plugin_lib: Path): + import tensorrt as trt + + self.trt = trt + self.logger = trt.Logger(trt.Logger.WARNING) + trt.init_libnvinfer_plugins(self.logger, "") + load_plugin_library(str(plugin_lib)) + if not find_plugin_creator(trt): + raise RuntimeError(f"{PLUGIN_NAME} plugin creator is not registered after loading {plugin_lib}") + + self.runtime = trt.Runtime(self.logger) + self.engine = self.runtime.deserialize_cuda_engine(engine_path.read_bytes()) + if self.engine is None: + raise RuntimeError(f"Failed to deserialize TensorRT engine: {engine_path}") + self.context = self.engine.create_execution_context() + if self.context is None: + raise RuntimeError(f"Failed to create execution context: {engine_path}") + + names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)] + missing = [name for name in ("left", "right", "disp") if name not in names] + if missing: + raise RuntimeError( + f"Plugin engine must expose tensors named left, right, and disp. " + f"Missing {missing}; found {names}" + ) + + def _torch_dtype(self, trt_dtype): + trt = self.trt + mapping = { + trt.DataType.FLOAT: torch.float32, + trt.DataType.HALF: torch.float16, + trt.DataType.INT32: torch.int32, + trt.DataType.INT8: torch.int8, + trt.DataType.BOOL: torch.bool, + } + if hasattr(trt.DataType, "BF16"): + mapping[trt.DataType.BF16] = torch.bfloat16 + if trt_dtype not in mapping: + raise RuntimeError(f"Unsupported TensorRT dtype: {trt_dtype}") + return mapping[trt_dtype] + + def infer(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + inputs = {"left": left, "right": right} + for name, tensor in list(inputs.items()): + expected = self._torch_dtype(self.engine.get_tensor_dtype(name)) + if tensor.dtype != expected: + tensor = tensor.to(expected) + inputs[name] = tensor.contiguous() + self.context.set_input_shape(name, tuple(inputs[name].shape)) + + out_shape = tuple(self.context.get_tensor_shape("disp")) + out_dtype = self._torch_dtype(self.engine.get_tensor_dtype("disp")) + output = torch.empty(out_shape, device="cuda", dtype=out_dtype) + + self.context.set_tensor_address("left", int(inputs["left"].data_ptr())) + self.context.set_tensor_address("right", int(inputs["right"].data_ptr())) + self.context.set_tensor_address("disp", int(output.data_ptr())) + + stream = torch.cuda.current_stream().cuda_stream + if not self.context.execute_async_v3(stream): + raise RuntimeError("TensorRT enqueue failed") + return output + + +def main() -> int: + args = parse_args() + set_logging_format() + set_seed(0) + torch.autograd.set_grad_enabled(False) + args.out_dir.mkdir(parents=True, exist_ok=True) + + engine_path = resolve_engine_path(args.model_dir, args.engine_file) + config_path = resolve_config_path(engine_path, args.model_dir, args.config_file) + plugin_lib = args.plugin_lib or find_default_plugin_library() + if plugin_lib is None or not plugin_lib.exists(): + raise FileNotFoundError( + "Could not find libffs_gwc_plugin.so. Build cpp first or pass " + "--plugin_lib /path/to/libffs_gwc_plugin.so" + ) + + with config_path.open("r") as f: + cfg = yaml.safe_load(f) + model_h, model_w = [int(v) for v in cfg["image_size"]] + + left = load_rgb(args.left_file) + right = load_rgb(args.right_file) + if left.shape != right.shape: + raise RuntimeError(f"left/right image size mismatch: {left.shape} vs {right.shape}") + input_h, input_w = left.shape[:2] + + logging.info(f"Engine: {engine_path}") + logging.info(f"Plugin: {plugin_lib}") + logging.info(f"Config: {config_path}") + logging.info(f"Input images: {input_w}x{input_h}") + logging.info(f"Model target resolution: {model_w}x{model_h}") + + left_model, scaled_h, scaled_w = resize_uniform_and_pad_rgb(left, model_h, model_w) + right_model, right_scaled_h, right_scaled_w = resize_uniform_and_pad_rgb(right, model_h, model_w) + if (scaled_h, scaled_w) != (right_scaled_h, right_scaled_w): + raise RuntimeError("left/right images produced different resize scales") + if (input_h, input_w) != (model_h, model_w): + logging.info( + f"Uniform resize+pad: {input_w}x{input_h} -> " + f"{scaled_w}x{scaled_h} inside {model_w}x{model_h}" + ) + + runner = PluginTensorRTRunner(engine_path, plugin_lib) + t_left = tensor_from_rgb_255(left_model) + t_right = tensor_from_rgb_255(right_model) + + logging.info("Running TensorRT inference") + disp_raw_t = runner.infer(t_left, t_right) + torch.cuda.current_stream().synchronize() + logging.info("Inference done") + + disp_raw = disp_raw_t.float().detach().cpu().numpy() + disp = postprocess_disparity( + disp_raw, + input_h=input_h, + input_w=input_w, + model_h=model_h, + model_w=model_w, + scaled_h=scaled_h, + scaled_w=scaled_w, + ) + + if not np.isfinite(disp).any(): + raise RuntimeError("Model produced no finite disparity values") + + k, baseline = load_intrinsics(args.intrinsic_file) + depth = disparity_to_depth_cpp(disp, float(k[0, 0]), baseline) + + disparity_path = args.out_dir / "disparity.bin" + depth_bin_path = args.out_dir / "depth_meter.bin" + depth_npy_path = args.out_dir / "depth_meter.npy" + disp_vis_path = args.out_dir / "disp_vis.png" + depth_vis_path = args.out_dir / "depth_vis.png" + + save_float_matrix_bin(disparity_path, disp) + save_float_matrix_bin(depth_bin_path, depth) + np.save(depth_npy_path, depth) + + disp_color = vis_disparity(disp, color_map=cv2.COLORMAP_TURBO) + disp_vis = np.concatenate([left, right, disp_color], axis=1) + imageio.imwrite(disp_vis_path, disp_vis) + imageio.imwrite(depth_vis_path, colorize_depth(depth)) + + logging.info(f"Saved: {disparity_path}") + logging.info(f"Saved: {depth_bin_path}") + logging.info(f"Saved: {depth_npy_path}") + logging.info(f"Saved: {disp_vis_path}") + logging.info(f"Saved: {depth_vis_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cv/fast-foundationstereo/scripts/run_demo_single_trt.py b/cv/fast-foundationstereo/scripts/run_demo_single_trt.py new file mode 100644 index 0000000..b78c749 --- /dev/null +++ b/cv/fast-foundationstereo/scripts/run_demo_single_trt.py @@ -0,0 +1,348 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Run Fast FoundationStereo inference with the single ONNX model (or TRT engine) +produced by make_single_onnx.py. + +Supports two backends: + - ONNX Runtime (default if --model_file points to an .onnx, or auto-detected) + - TensorRT (if --model_file points to an .engine) + +The model expects ImageNet-normalised inputs, so this script applies +normalisation during preprocessing. + +Usage: + # Run directly with ONNX (no trtexec step needed): + python run_demo_single_trt.py \ + --model_dir ./output_single_onnx \ + --left_file ../demo_data/left.png \ + --right_file ../demo_data/right.png + + # Or with an explicit model file: + python run_demo_single_trt.py \ + --model_dir ./output_single_onnx \ + --model_file ./output_single_onnx/fast_foundationstereo.onnx \ + --left_file ../demo_data/left.png \ + --right_file ../demo_data/right.png +""" + +import argparse +import logging +import os +import sys + +import cv2 +import imageio +import numpy as np +import torch +import yaml + +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') + +from Utils import ( + set_logging_format, set_seed, vis_disparity, + depth2xyzmap, toOpen3dCloud, o3d, +) + +IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) +IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) + + +class SingleEngineTrtRunner: + """Minimal TensorRT runner for a single engine with named I/O.""" + + def __init__(self, engine_path): + import tensorrt as trt + self.trt = trt + self.logger = trt.Logger(trt.Logger.WARNING) + + with open(engine_path, 'rb') as f: + self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read()) + if self.engine is None: + raise RuntimeError( + f'Failed to deserialize TRT engine from {engine_path}. ' + f'This usually means the engine was built with a different ' + f'TensorRT version (yours: {trt.__version__}). ' + f'Rebuild with: trtexec --onnx= ' + f'--saveEngine={engine_path} --fp16') + self.context = self.engine.create_execution_context() + + def _trt_to_torch_dtype(self, dt): + trt = self.trt + mapping = { + trt.DataType.FLOAT: torch.float32, + trt.DataType.HALF: torch.float16, + trt.DataType.BF16: torch.bfloat16, + trt.DataType.INT32: torch.int32, + trt.DataType.INT8: torch.int8, + trt.DataType.BOOL: torch.bool, + } + if dt not in mapping: + raise RuntimeError(f'Unsupported TRT dtype: {dt}') + return mapping[dt] + + def __call__(self, inputs: dict) -> dict: + """Run inference. + + Args: + inputs: {binding_name: torch.Tensor} for every input tensor. + Returns: + {binding_name: torch.Tensor} for every output tensor. + """ + trt = self.trt + + for name, tensor in inputs.items(): + expected = self._trt_to_torch_dtype(self.engine.get_tensor_dtype(name)) + if tensor.dtype != expected: + inputs[name] = tensor.to(expected) + if not inputs[name].is_contiguous(): + inputs[name] = inputs[name].contiguous() + self.context.set_input_shape(name, tuple(inputs[name].shape)) + + out_names = [ + self.engine.get_tensor_name(i) + for i in range(self.engine.num_io_tensors) + if self.engine.get_tensor_mode(self.engine.get_tensor_name(i)) + == trt.TensorIOMode.OUTPUT + ] + + outputs = {} + for name in out_names: + shape = tuple(self.context.get_tensor_shape(name)) + dtype = self._trt_to_torch_dtype(self.engine.get_tensor_dtype(name)) + outputs[name] = torch.empty(shape, device='cuda', dtype=dtype) + + for name, tensor in inputs.items(): + self.context.set_tensor_address(name, int(tensor.data_ptr())) + for name, tensor in outputs.items(): + self.context.set_tensor_address(name, int(tensor.data_ptr())) + + stream = torch.cuda.current_stream().cuda_stream + assert self.context.execute_async_v3(stream) + + return outputs + + +class OnnxRuntimeRunner: + """Run inference via ONNX Runtime (GPU if available, else CPU).""" + + def __init__(self, onnx_path): + import onnxruntime as ort + providers = [] + if 'CUDAExecutionProvider' in ort.get_available_providers(): + providers.append('CUDAExecutionProvider') + providers.append('CPUExecutionProvider') + logging.info(f'ONNX Runtime providers: {providers}') + self.session = ort.InferenceSession(onnx_path, providers=providers) + self.input_names = [inp.name for inp in self.session.get_inputs()] + self.output_names = [out.name for out in self.session.get_outputs()] + + def __call__(self, inputs: dict) -> dict: + feed = {} + for name in self.input_names: + tensor = inputs[name] + if isinstance(tensor, torch.Tensor): + tensor = tensor.cpu().float().numpy() + feed[name] = tensor + raw_outputs = self.session.run(self.output_names, feed) + outputs = {} + for name, arr in zip(self.output_names, raw_outputs): + outputs[name] = torch.as_tensor(arr).cuda() + return outputs + + +def normalize_imagenet(img_uint8: np.ndarray) -> np.ndarray: + """Apply ImageNet normalization: (img/255 - mean) / std.""" + return ((img_uint8.astype(np.float32) / 255.0) - IMAGENET_MEAN) / IMAGENET_STD + + +def resolve_config(model_path: str) -> str: + """Find the YAML config matching the model file, falling back to defaults.""" + model_dir = os.path.dirname(model_path) + base = os.path.splitext(os.path.basename(model_path))[0] + candidates = [ + os.path.join(model_dir, f'{base}.yaml'), + os.path.join(model_dir, 'config.yaml'), + os.path.join(model_dir, 'onnx.yaml'), + ] + for p in candidates: + if os.path.exists(p): + return p + raise FileNotFoundError( + f'No .yaml config found for {model_path}. ' + 'Run make_single_onnx.py first.') + + +def find_model(model_dir: str) -> str: + """Find an .engine or .onnx file in the directory (prefer .engine).""" + for ext in ('.engine', '.onnx'): + for f in os.listdir(model_dir): + if f.endswith(ext): + return os.path.join(model_dir, f) + raise FileNotFoundError( + f'No .engine or .onnx file found in {model_dir}. ' + 'Run make_single_onnx.py first.') + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description='Run Fast FoundationStereo with ONNX Runtime or TensorRT') + parser.add_argument('--model_dir', type=str, + default=f'{code_dir}/output_single_onnx', + help='Directory containing .onnx/.engine + config.yaml') + parser.add_argument('--model_file', type=str, default='', + help='Explicit path to .onnx or .engine file (overrides auto-search)') + parser.add_argument('--left_file', type=str, + default=f'{code_dir}/../demo_data/left.png') + parser.add_argument('--right_file', type=str, + default=f'{code_dir}/../demo_data/right.png') + parser.add_argument('--intrinsic_file', type=str, + default=f'{code_dir}/../demo_data/K.txt', + help='Camera intrinsic matrix and baseline file') + parser.add_argument('--out_dir', type=str, + default=f'{code_dir}/../output_demo') + parser.add_argument('--remove_invisible', type=int, default=1) + parser.add_argument('--denoise_cloud', type=int, default=1) + parser.add_argument('--denoise_nb_points', type=int, default=30) + parser.add_argument('--denoise_radius', type=float, default=0.03) + parser.add_argument('--get_pc', type=int, default=1, + help='Generate and save point cloud') + parser.add_argument('--zfar', type=float, default=100, + help='Max depth (m) to include in point cloud') + args = parser.parse_args() + + set_logging_format() + set_seed(0) + torch.autograd.set_grad_enabled(False) + os.makedirs(args.out_dir, exist_ok=True) + + # ── Find model and config ───────────────────────────────────────────── + model_path = args.model_file if args.model_file else find_model(args.model_dir) + cfg_path = resolve_config(model_path) + with open(cfg_path, 'r') as f: + cfg = yaml.safe_load(f) + target_h, target_w = cfg['image_size'] + logging.info(f'Model target resolution: {target_h} x {target_w}') + + # ── Load model (ONNX Runtime or TensorRT) ──────────────────────────── + logging.info(f'Loading model: {model_path}') + if model_path.endswith('.onnx'): + runner = OnnxRuntimeRunner(model_path) + else: + runner = SingleEngineTrtRunner(model_path) + + # ── Read images ─────────────────────────────────────────────────────── + img0 = imageio.imread(args.left_file) + img1 = imageio.imread(args.right_file) + + if img0.ndim == 2: + img0 = np.tile(img0[..., None], (1, 1, 3)) + img1 = np.tile(img1[..., None], (1, 1, 3)) + img0 = img0[..., :3] + img1 = img1[..., :3] + + # ── Resize to model resolution (direct stretch) ──────────────────────── + orig_h, orig_w = img0.shape[:2] + fx = target_w / orig_w + fy = target_h / orig_h + + if fx != 1 or fy != 1: + logging.info( + f'Resizing images: {orig_h}x{orig_w} → {target_h}x{target_w} ' + f'(fx={fx:.4f}, fy={fy:.4f})') + img0 = cv2.resize(img0, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + img1 = cv2.resize(img1, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + H, W = img0.shape[:2] + + img0_ori = img0.copy() + img1_ori = img1.copy() + logging.info(f'Image size after resize: {img0.shape}') + imageio.imwrite(f'{args.out_dir}/left.png', img0) + imageio.imwrite(f'{args.out_dir}/right.png', img1) + + # ── Preprocess: ImageNet normalize → NCHW float tensor ──────────────── + img0_norm = normalize_imagenet(img0) + img1_norm = normalize_imagenet(img1) + + t_left = torch.as_tensor(img0_norm).cuda().float()[None].permute(0, 3, 1, 2) + t_right = torch.as_tensor(img1_norm).cuda().float()[None].permute(0, 3, 1, 2) + + # ── Inference ───────────────────────────────────────────────────────── + logging.info('Running inference (first run may be slow due to TRT warmup)') + outputs = runner({'left_image': t_left, 'right_image': t_right}) + disp = outputs['disparity'] + logging.info('Inference done') + + disp = disp.float().cpu().numpy().reshape(H, W).clip(0, None) * (1.0 / fx) + + # ── Visualise disparity ────────────────────────────────────────────── + vis = vis_disparity(disp, color_map=cv2.COLORMAP_TURBO) + vis = np.concatenate([img0_ori, img1_ori, vis], axis=1) + imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis) + s = 1280 / vis.shape[1] + resized_vis = cv2.resize(vis, (int(vis.shape[1] * s), int(vis.shape[0] * s))) + cv2.imshow('disp', resized_vis[:, :, ::-1]) + cv2.waitKey(0) + + # ── Remove invisible pixels ────────────────────────────────────────── + if args.remove_invisible: + _, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij') + invalid = (xx - disp) < 0 + disp[invalid] = np.inf + + # ── Point cloud generation ─────────────────────────────────────────── + if args.get_pc: + with open(args.intrinsic_file, 'r') as f: + lines = f.readlines() + K = (np.array(list(map(float, lines[0].rstrip().split()))) + .astype(np.float32).reshape(3, 3)) + baseline = float(lines[1]) + K[:2] *= np.array([fx, fy], dtype=np.float32)[:, np.newaxis] + depth = K[0, 0] * baseline / disp + np.save(f'{args.out_dir}/depth_meter.npy', depth) + + xyz_map = depth2xyzmap(depth, K) + pcd = toOpen3dCloud(xyz_map.reshape(-1, 3), img0_ori.reshape(-1, 3)) + pts = np.asarray(pcd.points) + keep = (pts[:, 2] > 0) & (pts[:, 2] <= args.zfar) + pcd = pcd.select_by_index(np.where(keep)[0]) + o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd) + logging.info(f'Point cloud saved to {args.out_dir}') + + if args.denoise_cloud: + logging.info('Denoising point cloud...') + _, ind = pcd.remove_radius_outlier( + nb_points=args.denoise_nb_points, + radius=args.denoise_radius) + pcd = pcd.select_by_index(ind) + o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', pcd) + + logging.info('Visualizing point cloud. Press ESC to exit.') + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.add_geometry(pcd) + vis.get_render_option().point_size = 1.0 + vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5]) + ctr = vis.get_view_control() + ctr.set_front([0, 0, -1]) + closest = np.asarray(pcd.points)[:, 2].argmin() + ctr.set_lookat(np.asarray(pcd.points)[closest]) + ctr.set_up([0, -1, 0]) + vis.run() + vis.destroy_window() diff --git a/cv/fast-foundationstereo/scripts/run_demo_tensorrt.py b/cv/fast-foundationstereo/scripts/run_demo_tensorrt.py new file mode 100755 index 0000000..9f8244d --- /dev/null +++ b/cv/fast-foundationstereo/scripts/run_demo_tensorrt.py @@ -0,0 +1,145 @@ +import os,sys +code_dir = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(f'{code_dir}/../') +from omegaconf import OmegaConf +from core.utils.utils import InputPadder +import argparse, torch, logging, yaml +import imageio +import numpy as np +from Utils import ( + set_logging_format, set_seed, vis_disparity, + depth2xyzmap, toOpen3dCloud, o3d, +) +from core.foundation_stereo import TrtRunner +import cv2 + + +def resolve_onnx_cfg_path(onnx_dir: str) -> str: + onnx_dir = os.path.normpath(onnx_dir) + candidates = [ + os.path.join(onnx_dir, 'onnx.yaml'), + os.path.join(os.path.dirname(onnx_dir), 'onnx.yaml'), + ] + for p in candidates: + if os.path.exists(p): + return p + raise FileNotFoundError( + f"onnx.yaml not found. Looked in: {candidates}. " + "Please run scripts/make_onnx.py first to generate ONNX metadata." + ) + + +if __name__=="__main__": + code_dir = os.path.dirname(os.path.realpath(__file__)) + parser = argparse.ArgumentParser() + parser.add_argument('--onnx_dir', default=f'{code_dir}/output', type=str) + parser.add_argument('--left_file', default=f'{code_dir}/../assets/left.png', type=str) + parser.add_argument('--right_file', default=f'{code_dir}/../assets/right.png', type=str) + parser.add_argument('--intrinsic_file', default=f'{code_dir}/../assets/K.txt', type=str, help='camera intrinsic matrix and baseline file') + parser.add_argument('--out_dir', default='/home/bowen/debug/stereo_output', type=str) + parser.add_argument('--remove_invisible', default=1, type=int) + parser.add_argument('--denoise_cloud', default=1, type=int) + parser.add_argument('--denoise_nb_points', type=int, default=30, help='number of points to consider for radius outlier removal') + parser.add_argument('--denoise_radius', type=float, default=0.03, help='radius to use for outlier removal') + parser.add_argument('--get_pc', type=int, default=1, help='save point cloud output') + parser.add_argument('--zfar', type=float, default=100, help="max depth to include in point cloud") + args = parser.parse_args() + + set_logging_format() + set_seed(0) + torch.autograd.set_grad_enabled(False) + os.makedirs(args.out_dir, exist_ok=True) + + onnx_cfg_path = resolve_onnx_cfg_path(args.onnx_dir) + with open(onnx_cfg_path, 'r') as ff: + cfg:dict = yaml.safe_load(ff) + for k in args.__dict__: + if args.__dict__[k] is not None: + cfg[k] = args.__dict__[k] + args = OmegaConf.create(cfg) + logging.info(f"args:\n{args}") + model = TrtRunner(args, args.onnx_dir+'/feature_runner.engine', args.onnx_dir+'/post_runner.engine') + + img0 = imageio.imread(args.left_file) + img1 = imageio.imread(args.right_file) + if len(img0.shape)==2: + img0 = np.tile(img0[...,None], (1,1,3)) + img1 = np.tile(img1[...,None], (1,1,3)) + img0 = img0[...,:3] + img1 = img1[...,:3] + H,W = img0.shape[:2] + + fx = args.image_size[1] / img0.shape[1] + fy = args.image_size[0] / img0.shape[0] + if fx != 1 or fy != 1: + logging.info(f">>>>>>>>>>>>>>>WARNING: resizing image to {args.image_size}, fx: {fx}, fy: {fy}, this is not recommended. It's best to make tensorrt engine with the same image size as the input image.") + img0 = cv2.resize(img0, fx=fx, fy=fy, dsize=None) + img1 = cv2.resize(img1, fx=fx, fy=fy, dsize=None) + H,W = img0.shape[:2] + img0_ori = img0.copy() + img1_ori = img1.copy() + logging.info(f"img0: {img0.shape}") + imageio.imwrite(f'{args.out_dir}/left.png', img0) + imageio.imwrite(f'{args.out_dir}/right.png', img1) + + img0 = torch.as_tensor(img0).cuda().float()[None].permute(0,3,1,2) + img1 = torch.as_tensor(img1).cuda().float()[None].permute(0,3,1,2) + + logging.info(f"Start forward, 1st time run can be slow due to compilation") + disp = model.forward(img0, img1) + logging.info("forward done") + disp = disp.data.cpu().numpy().reshape(H,W).clip(0, None) * 1/fx + + cmap = None + min_val = None + max_val = None + vis = vis_disparity(disp, min_val=min_val, max_val=max_val, cmap=cmap, color_map=cv2.COLORMAP_TURBO) + vis = np.concatenate([img0_ori, img1_ori, vis], axis=1) + imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis) + s = 1280/vis.shape[1] + resized_vis = cv2.resize(vis, (int(vis.shape[1]*s), int(vis.shape[0]*s))) + cv2.imshow('disp', resized_vis[:,:,::-1]) + cv2.waitKey(0) + + if args.remove_invisible: + yy,xx = np.meshgrid(np.arange(disp.shape[0]), np.arange(disp.shape[1]), indexing='ij') + us_right = xx-disp + invalid = us_right<0 + disp[invalid] = np.inf + + if args.get_pc: + with open(args.intrinsic_file, 'r') as f: + lines = f.readlines() + K = np.array(list(map(float, lines[0].rstrip().split()))).astype(np.float32).reshape(3,3) + baseline = float(lines[1]) + K[:2] *= np.array([fx, fy], dtype=np.float32)[:, np.newaxis] + depth = K[0,0]*baseline/disp + np.save(f'{args.out_dir}/depth_meter.npy', depth) + xyz_map = depth2xyzmap(depth, K) + pcd = toOpen3dCloud(xyz_map.reshape(-1,3), img0_ori.reshape(-1,3)) + keep_mask = (np.asarray(pcd.points)[:,2]>0) & (np.asarray(pcd.points)[:,2]<=args.zfar) + keep_ids = np.arange(len(np.asarray(pcd.points)))[keep_mask] + pcd = pcd.select_by_index(keep_ids) + o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd) + logging.info(f"PCL saved to {args.out_dir}") + + if args.denoise_cloud: + logging.info("[Optional step] denoise point cloud...") + cl, ind = pcd.remove_radius_outlier(nb_points=args.denoise_nb_points, radius=args.denoise_radius) + inlier_cloud = pcd.select_by_index(ind) + o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', inlier_cloud) + pcd = inlier_cloud + + logging.info("Visualizing point cloud. Press ESC to exit.") + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.add_geometry(pcd) + vis.get_render_option().point_size = 1.0 + vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5]) + ctr = vis.get_view_control() + ctr.set_front([0, 0, -1]) + id = np.asarray(pcd.points)[:,2].argmin() + ctr.set_lookat(np.asarray(pcd.points)[id]) + ctr.set_up([0, -1, 0]) + vis.run() + vis.destroy_window() diff --git a/cv/fast-foundationstereo/weights/.gitkeep b/cv/fast-foundationstereo/weights/.gitkeep new file mode 100755 index 0000000..e69de29