use bytemuck::{Pod,Zeroable}; use std::borrow::Cow; use wgpu::util::DeviceExt; const W:usize=832; const H:usize=480; #[repr(C)] #[derive(Copy,Clone,Pod,Zeroable)] struct Inst{c:[f32;2],a1:[f32;2],a2:[f32;2],pad:[f32;2],col:[f32;4]} const WGSL:&str=r#" struct Inst{c:vec2f,a1:vec2f,a2:vec2f,pad:vec2f,col:vec4f}; @group(0)@binding(0) var insts:array; struct VOut{@builtin(position) p:vec4f,@location(0) l:vec2f,@location(1) c:vec4f}; @vertex fn vs(@builtin(vertex_index) vi:u32,@builtin(instance_index) ii:u32)->VOut{ var cs=array(vec2f(-1.,-1.),vec2f(1.,-1.),vec2f(-1.,1.),vec2f(-1.,1.),vec2f(1.,-1.),vec2f(1.,1.)); let q=cs[vi]; let ins=insts[ii]; let p=ins.c+q.x*ins.a1+q.y*ins.a2; var o:VOut; o.p=vec4f(p.x/416.-1.,1.-p.y/240.,0.5,1.); o.l=q; o.c=ins.col; return o; } @fragment fn fs(i:VOut)->@location(0) vec4f{ let a=i.c.a*exp(-4.5*dot(i.l,i.l)); if(a<0.003){discard;} return vec4f(i.c.rgb,a); } "#; fn main(){ let args:Vec=std::env::args().collect(); let path=args.get(1).map(|s|s.as_str()).unwrap_or(r"D:\3DWord\itseez3d\ai-lab\ml-sharp\output\photo.ply"); let raw=std::fs::read(path).unwrap(); let hs=find(&raw,b"end_header ").unwrap()+11; let head=String::from_utf8_lossy(&raw[..hs]); let n:usize=head.lines().find(|l|l.starts_with("element vertex")).unwrap().split_whitespace().nth(2).unwrap().parse().unwrap(); let floats:Vec=raw[hs..hs+n*14*4].chunks_exact(4).map(|b| f32::from_le_bytes([b[0],b[1],b[2],b[3]])).collect(); let mut ps=Vec::with_capacity(n); for i in 0..n{let o=i*14; ps.push([floats[o],floats[o+1],floats[o+2]]);} let (mut mn,mut mx)=([f32::MAX;3],[f32::MIN;3]); for p in &ps{for k in 0..3{mn[k]=mn[k].min(p[k]);mx[k]=mx[k].max(p[k]);}} let ctr=[(mn[0]+mx[0])/2.,(mn[1]+mx[1])/2.,(mn[2]+mx[2])/2.]; let mut rad=0.0f32; for p in &ps{let d=(p[0]-ctr[0]).powi(2)+(p[1]-ctr[1]).powi(2)+(p[2]-ctr[2]).powi(2);rad=rad.max(d);} let rad=rad.sqrt(); let yaw=args.get(2).and_then(|s|s.parse().ok()).unwrap_or(0.0f32); let (cy,sy)=(yaw.cos(),yaw.sin()); let orbit=args.get(5).map(|s|s=="orbit").unwrap_or(false); let dist=1.8*rad; let eye=if orbit{[ctr[0]+dist*sy,ctr[1]+0.15*rad,ctr[2]+dist*cy]}else{[0.,0.,0.]}; let mut rv=[[1.,0.,0.],[0.,1.,0.],[0.,0.,1.]]; if orbit{let ex=ctr[0]-eye[0];let ey=ctr[1]-eye[1];let ez=ctr[2]-eye[2];let el=(ex*ex+ey*ey+ez*ez).sqrt();let f=[ex/el,ey/el,ez/el];let rx=-f[2];let rz=f[0];let rl=(rx*rx+rz*rz).sqrt();let r=[rx/rl,0.,rz/rl];let u=[f[1]*r[2]-f[2]*r[1],f[2]*r[0]-f[0]*r[2],f[0]*r[1]-f[1]*r[0]];rv=[r,u,f];} let covmin=args.get(4).and_then(|s|s.parse().ok()).unwrap_or(0.10f32); let fx=416.0f32; let fy=416.0f32; let cx0=416.0f32; let cy0=240.0f32; let mut tmp:Vec<(f32,Inst)>=Vec::with_capacity(n); for i in 0..n{ let o=i*14; let dc=[floats[o+3],floats[o+4],floats[o+5]]; let alpha=(1.0/(1.0+(-floats[o+6]).exp())*2.0).min(1.0); let sc=[floats[o+7].exp(),floats[o+8].exp(),floats[o+9].exp()]; let (mut w,mut x,mut y,mut z)=(floats[o+10],floats[o+11],floats[o+12],floats[o+13]); let qn=(w*w+x*x+y*y+z*z).sqrt(); if qn<1e-8{w=1.;x=0.;y=0.;z=0.;}else{w/=qn;x/=qn;y/=qn;z/=qn;} let r=[[1.-2.*(y*y+z*z),2.*(x*y-w*z),2.*(x*z+w*y)],[2.*(x*y+w*z),1.-2.*(x*x+z*z),2.*(y*z-w*x)],[2.*(x*z-w*y),2.*(y*z+w*x),1.-2.*(x*x+y*y)]]; let mut cov=[[0f32;3];3]; for a in 0..3{for b in 0..3{let mut t=0.;for k in 0..3{t+=r[a][k]*sc[k]*sc[k]*r[b][k];}cov[a][b]=t;}} let xc=cy*ps[i][0]-sy*ps[i][2]; let yc=ps[i][1]; let zc=sy*ps[i][0]+cy*ps[i][2]; if zc<0.05{continue;} let u=fx*xc/zc+cx0; let v=fy*yc/zc+cy0; if u < -50.||u>W as f32+50.||v < -50.||v>H as f32+50.{continue;} let cxx=cov[0][0]; let cxz=cov[0][2]; let czz=cov[2][2]; let r00=cy*cy*cxx-2.*cy*sy*cxz+sy*sy*czz; let r02=cy*sy*(cxx-czz)+(cy*cy-sy*sy)*cxz; let r22=sy*sy*cxx+2.*sy*cy*cxz+cy*cy*czz; let r01=cy*cov[0][1]-sy*cov[2][1]; let r12=sy*cov[0][1]+cy*cov[2][1]; let iz=1./zc; let j00=fx*iz; let j02=-fx*xc*iz*iz; let j11=fy*iz; let j12=-fy*yc*iz*iz; let a2=j00*j00*r00+2.*j00*j02*r02+j02*j02*r22; let b2=j00*j11*r01+j00*j12*r02+j02*j11*r12+j02*j12*r22; let c2=j11*j11*cov[1][1]+2.*j11*j12*r12+j12*j12*r22; let mid=(a2+c2)/2.; let d=(((a2-c2)/2.).powi(2)+b2*b2).sqrt(); let l1=(mid+d).max(0.01);let l2=(mid-d).max(0.01); let (mut v1x,mut v1y)=(b2,l1-a2); if v1x.abs()+v1y.abs()<1e-6{v1x=1.;v1y=0.;} let nn=(v1x*v1x+v1y*v1y).sqrt();v1x/=nn;v1y/=nn; let v2x=-v1y;let v2y=v1x; let s1=((3.*l1).sqrt()*2.0).clamp(covmin,12.); let s2=((3.*l2).sqrt()*2.0).clamp(covmin,12.); if !s1.is_finite()||!s2.is_finite()||!u.is_finite()||!v.is_finite()||alpha<0.003{continue;} let col=[dc[0].clamp(0.,1.),dc[1].clamp(0.,1.),dc[2].clamp(0.,1.)]; tmp.push((zc,Inst{c:[u,v],a1:[v1x*s1,v1y*s1],a2:[v2x*s2,v2y*s2],pad:[0.,0.],col:[col[0],col[1],col[2],alpha]})); } tmp.sort_by(|a,b|b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); let inst:Vec=tmp.into_iter().map(|x|x.1).collect(); println!("gaussians={} visible={}",n,inst.len()); let instance=wgpu::Instance::new(&wgpu::InstanceDescriptor{backends:wgpu::Backends::VULKAN,..Default::default()}); let ad=pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions{power_preference:wgpu::PowerPreference::HighPerformance,compatible_surface:None,force_fallback_adapter:false})).unwrap(); let info=ad.get_info();println!("{} {:?}",info.name,info.backend); let mut lim=wgpu::Limits::downlevel_defaults();lim.max_storage_buffer_binding_size=512<<20; let (dev,q)=pollster::block_on(ad.request_device(&wgpu::DeviceDescriptor{label:None,required_features:wgpu::Features::empty(),required_limits:lim,..Default::default()})).unwrap(); let ib=dev.create_buffer_init(&wgpu::util::BufferInitDescriptor{label:None,contents:bytemuck::cast_slice(&inst),usage:wgpu::BufferUsages::STORAGE}); let tex=dev.create_texture(&wgpu::TextureDescriptor{label:None,size:wgpu::Extent3d{width:W as u32,height:H as u32,depth_or_array_layers:1},mip_level_count:1,sample_count:1,dimension:wgpu::TextureDimension::D2,format:wgpu::TextureFormat::Rgba8UnormSrgb,usage:wgpu::TextureUsages::RENDER_ATTACHMENT|wgpu::TextureUsages::COPY_SRC,view_formats:&[]}); let view=tex.create_view(&wgpu::TextureViewDescriptor::default()); let md=dev.create_shader_module(wgpu::ShaderModuleDescriptor{label:None,source:wgpu::ShaderSource::Wgsl(Cow::Borrowed(WGSL))}); let pl=dev.create_render_pipeline(&wgpu::RenderPipelineDescriptor{label:None,layout:None,vertex:wgpu::VertexState{module:&md,entry_point:Some("vs"),compilation_options:Default::default(),buffers:&[]},fragment:Some(wgpu::FragmentState{module:&md,entry_point:Some("fs"),compilation_options:Default::default(),targets:&[Some(wgpu::ColorTargetState{format:wgpu::TextureFormat::Rgba8UnormSrgb,blend:Some(wgpu::BlendState::ALPHA_BLENDING),write_mask:wgpu::ColorWrites::ALL})]}),primitive:wgpu::PrimitiveState::default(),depth_stencil:None,multisample:wgpu::MultisampleState::default(),multiview:None,cache:None}); let bgl=pl.get_bind_group_layout(0); let bg=dev.create_bind_group(&wgpu::BindGroupDescriptor{label:None,layout:&bgl,entries:&[wgpu::BindGroupEntry{binding:0,resource:ib.as_entire_binding()}]}); let mut en=dev.create_command_encoder(&wgpu::CommandEncoderDescriptor{label:None}); {let mut rp=en.begin_render_pass(&wgpu::RenderPassDescriptor{label:None,color_attachments:&[Some(wgpu::RenderPassColorAttachment{view:&view,resolve_target:None,depth_slice:None,ops:wgpu::Operations{load:wgpu::LoadOp::Clear(wgpu::Color{r:0.55,g:0.72,b:0.9,a:1.}),store:wgpu::StoreOp::Store}})],depth_stencil_attachment:None,timestamp_writes:None,occlusion_query_set:None});rp.set_pipeline(&pl);rp.set_bind_group(0,&bg,&[]);rp.draw(0..6,0..inst.len() as u32);} let bpr=(W*4) as u32; let rb=dev.create_buffer(&wgpu::BufferDescriptor{label:None,size:(bpr*H as u32) as u64,usage:wgpu::BufferUsages::COPY_DST|wgpu::BufferUsages::MAP_READ,mapped_at_creation:false}); en.copy_texture_to_buffer(wgpu::TexelCopyTextureInfo{texture:&tex,mip_level:0,origin:wgpu::Origin3d::ZERO,aspect:wgpu::TextureAspect::All},wgpu::TexelCopyBufferInfo{buffer:&rb,layout:wgpu::TexelCopyBufferLayout{offset:0,bytes_per_row:Some(bpr),rows_per_image:Some(H as u32)}},wgpu::Extent3d{width:W as u32,height:H as u32,depth_or_array_layers:1}); q.submit(Some(en.finish()));let sl=rb.slice(..);sl.map_async(wgpu::MapMode::Read,|r|r.unwrap());dev.poll(wgpu::PollType::wait_indefinitely()).unwrap();let d=sl.get_mapped_range();let mut img=Vec::with_capacity(W*H*3);for y in 0..H{for x in 0..W{let o=y*bpr as usize+x*4;img.push(d[o]);img.push(d[o+1]);img.push(d[o+2]);}} let outp=args.get(3).map(|s|s.to_string()).unwrap_or(r"D:\3DWord\itseez3d\ai-lab\gaussian-splatting\vulkan-rs\splat.png".to_string()); png_rgb(&outp,W as u32,H as u32,&img); println!("saved {}",outp); } fn find(h:&[u8],n:&[u8])->Option{(0..h.len()).find(|i|h.get(*i..*i+n.len())==Some(n))} fn crc32(b:&[u8])->u32{let mut c=0xffffffffu32;for &x in b{c^=x as u32;for _ in 0..8{c=if c&1!=0{0xedb88320^(c>>1)}else{c>>1};}}c^0xffffffff} fn adler32(b:&[u8])->u32{let (mut a,mut d)=(1u32,0u32);for &x in b{a=(a+x as u32)%65521;d=(d+a)%65521;}(d<<16)|a} fn chunk(out:&mut Vec,typ:&[u8;4],data:&[u8]){out.extend_from_slice(&(data.len() as u32).to_be_bytes());let mut t=Vec::with_capacity(4+data.len());t.extend_from_slice(typ);t.extend_from_slice(data);out.extend_from_slice(&t);out.extend_from_slice(&crc32(&t).to_be_bytes());} fn png_rgb(path:&str,w:u32,h:u32,rgb:&[u8]){ let mut raw=Vec::with_capacity((w as usize*3+1)*h as usize); for y in 0..h as usize{raw.push(0);raw.extend_from_slice(&rgb[y*w as usize*3..(y+1)*w as usize*3]);} let mut z=vec![0x78,0x01]; let mut off=0; while off=raw.len();z.push(if last{1}else{0});z.extend_from_slice(&(n as u16).to_le_bytes());z.extend_from_slice(&(!(n as u16)).to_le_bytes());z.extend_from_slice(&raw[off..off+n]);off+=n;} z.extend_from_slice(&adler32(&raw).to_be_bytes()); let mut out=vec![137,80,78,71,13,10,26,10]; let mut ihdr=Vec::new();ihdr.extend_from_slice(&w.to_be_bytes());ihdr.extend_from_slice(&h.to_be_bytes());ihdr.extend_from_slice(&[8,2,0,0,0]);chunk(&mut out,b"IHDR",&ihdr);chunk(&mut out,b"IDAT",&z);chunk(&mut out,b"IEND",&[]);std::fs::write(path,out).unwrap(); }