1use pk_stl::parse_stl;
2use obj::{load_obj, Obj, TexturedVertex};
3use image::{Rgb, RgbImage, DynamicImage, ImageError};
4use rayon::prelude::*;
5use crate::spacial::tri::Tri;
6use crate::datatypes::material::Material;
7use crate::datatypes::vector3::Vector3;
8use crate::datatypes::vector2::Vector2;
9use crate::datatypes::color::Color;
10use crate::datatypes::vector2d::Vector2D;
11use crate::utilities::postprocessing::remove_fireflies;
12use std::fs;
13use std::sync::atomic::{AtomicUsize, Ordering};
14use std::sync::Arc;
15use std::fs::File;
16use std::io::BufReader;
17use std::path::Path;
18
19pub fn load_model(file_path: &str, material: Material) -> Vec<Tri> {
20 if file_path.ends_with(".obj") {
21 println!("Processing .obj file: {}", file_path);
22 return import_obj(file_path, material)
23 } else if file_path.ends_with(".stl") {
24 println!("Processing .stl file: {}", file_path);
25 return import_stl(file_path, material)
26 } else {
27 println!("Unsupported file extension");
28 return vec![]
29 }
30}
31
32fn import_stl(file_path: &str, material: Material) -> Vec<Tri> {
33 let content = fs::read(file_path).expect("Failed to read model file");
34 let model = parse_stl(content.as_slice()).unwrap();
35
36 println!("Computing vertex normals for {}", file_path);
37 let counter: Arc<AtomicUsize> = Arc::new(AtomicUsize::new(0));
38 let model_tris: Vec<Tri> = model.triangles.par_iter().map(|tri| {
39 let vertex1: Vector3 = Vector3::new(tri.vertices[0].x.into(), tri.vertices[0].y.into(), tri.vertices[0].z.into());
40 let vertex2: Vector3 = Vector3::new(tri.vertices[1].x.into(), tri.vertices[1].y.into(), tri.vertices[1].z.into());
41 let vertex3: Vector3 = Vector3::new(tri.vertices[2].x.into(), tri.vertices[2].y.into(), tri.vertices[2].z.into());
42 let face_normal = Vector3::new(tri.normal.x.into(), tri.normal.y.into(), tri.normal.z.into());
43
44 let mut vertex1_normal: Vector3 = Vector3::zero();
45 let mut vertex2_normal: Vector3 = Vector3::zero();
46 let mut vertex3_normal: Vector3 = Vector3::zero();
47
48 let e: f64 = 0.0001;
49 for other_tri in &model.triangles {
50 let other_vertex1: Vector3 = Vector3::new(other_tri.vertices[0].x.into(), other_tri.vertices[0].y.into(), other_tri.vertices[0].z.into());
51 let other_vertex2: Vector3 = Vector3::new(other_tri.vertices[1].x.into(), other_tri.vertices[1].y.into(), other_tri.vertices[1].z.into());
52 let other_vertex3: Vector3 = Vector3::new(other_tri.vertices[2].x.into(), other_tri.vertices[2].y.into(), other_tri.vertices[2].z.into());
53 let other_face_normal = Vector3::new(other_tri.normal.x.into(), other_tri.normal.y.into(), other_tri.normal.z.into());
54
55 if (vertex1 - other_vertex1).magnitude() < e
56 || (vertex1 - other_vertex2).magnitude() < e
57 || (vertex1 - other_vertex3).magnitude() < e
58 {
59 vertex1_normal += other_face_normal;
60 }
61 if (vertex2 - other_vertex1).magnitude() < e
62 || (vertex2 - other_vertex2).magnitude() < e
63 || (vertex2 - other_vertex3).magnitude() < e
64 {
65 vertex2_normal += other_face_normal;
66 }
67 if (vertex3 - other_vertex1).magnitude() < e
68 || (vertex3 - other_vertex2).magnitude() < e
69 || (vertex3 - other_vertex3).magnitude() < e
70 {
71 vertex3_normal += other_face_normal;
72 }
73 }
74
75 vertex1_normal = vertex1_normal.normalize();
76 vertex2_normal = vertex2_normal.normalize();
77 vertex3_normal = vertex3_normal.normalize();
78
79 let tri: Tri = Tri::new(
80 vertex1, vertex2, vertex3, vertex1_normal,
81 vertex2_normal, vertex3_normal, face_normal,
82 Vector2::zero(), Vector2::zero(),
83 Vector2::zero(), material
84 );
85
86 let current: usize = counter.fetch_add(1, Ordering::Relaxed);
87 if current % 1000 == 0 || model.triangles.len() == current {
88 println!("Computing vertex normals: {}/{}", current, model.triangles.len());
89 }
90 tri
91 }).collect();
92 println!("Done");
93 model_tris
94}
95
96fn import_obj(file_path: &str, material: Material) -> Vec<Tri> {
97 let input = BufReader::new(File::open(file_path).expect("Failed to read OBJ"));
98 let dome: Obj<TexturedVertex> = load_obj(input).expect("Failed to load OBJ");
99 let vertices = dome.vertices;
100 let indices = dome.indices;
101 let mut triangles: Vec<Tri> = Vec::new();
102
103 for chunk in indices.chunks(3) {
104 if chunk.len() != 3 {
105 continue;
106 }
107
108 let p1_index = chunk[0] as usize;
109 let p2_index = chunk[1] as usize;
110 let p3_index = chunk[2] as usize;
111
112 let vert1 = vertices[p1_index];
113 let vert2 = vertices[p2_index];
114 let vert3 = vertices[p3_index];
115
116 let p1: Vector3 = Vector3::new(vert1.position[0] as f64, vert1.position[2] as f64, vert1.position[1] as f64);
117 let p2: Vector3 = Vector3::new(vert2.position[0] as f64, vert2.position[2] as f64, vert2.position[1] as f64);
118 let p3: Vector3 = Vector3::new(vert3.position[0] as f64, vert3.position[2] as f64, vert3.position[1] as f64);
119
120 let p1_normal: Vector3 = Vector3::new(vert1.normal[0] as f64, vert1.normal[2] as f64, vert1.normal[1] as f64);
121 let p2_normal: Vector3 = Vector3::new(vert2.normal[0] as f64, vert2.normal[2] as f64, vert2.normal[1] as f64);
122 let p3_normal: Vector3 = Vector3::new(vert3.normal[0] as f64, vert3.normal[2] as f64, vert3.normal[1] as f64);
123
124 let p1_texture: Vector2 = Vector2::new(vert1.texture[0] as f64, vert1.texture[1] as f64);
125 let p2_texture: Vector2 = Vector2::new(vert2.texture[0] as f64, vert2.texture[1] as f64);
126 let p3_texture: Vector2 = Vector2::new(vert3.texture[0] as f64, vert3.texture[1] as f64);
127
128
129 let triangle = Tri::new(p1, p2, p3, p1_normal, p2_normal, p3_normal,
130 Tri::compute_face_normal(p1, p2, p3),
131 p1_texture, p2_texture, p3_texture, material
132 );
133
134 triangles.push(triangle);
135 }
136
137
138 triangles
139}
140
141pub fn import_texture(path: &str) -> Vector2D<Color> {
142 println!("Loading texture: {}", path);
143 let image_result: Result<DynamicImage, ImageError> = image::open(&Path::new(path));
144 if let Err(err) = image_result {
145 eprintln!("Failed to open image: {}", err);
146 return Vector2D::new(0, 0, Color::black());
147 }
148
149 let image: DynamicImage = image_result.unwrap();
150 let rgb_image: RgbImage = image.into_rgb8();
151
152 let width: usize = rgb_image.width() as usize;
153 let height: usize = rgb_image.height() as usize;
154 let mut pixel_vector: Vector2D<Color> = Vector2D::new(width, height, Color::black());
155
156 for (x, y, pixel) in rgb_image.enumerate_pixels() {
157 let color: Color = Color::new(pixel[0] as f64 / 255.0, pixel[1] as f64 / 255.0, pixel[2] as f64 / 255.0);
158 pixel_vector.set(y as usize, x as usize, color);
159 }
160
161 pixel_vector
162}
163
164pub fn save_vector2d_as_png(vector: &Vector2D<Color>, filename: &str) -> Result<(), image::ImageError> {
165 let mut image = RgbImage::new(vector.width as u32, vector.height as u32);
166 for (i, color) in remove_fireflies(vector).data.iter().enumerate() {
167 let x = (i % vector.width) as u32;
168 let y = (i / vector.width) as u32;
169
170 let rgb_color = Rgb([
171 (color.red * 255.0) as u8,
172 (color.green * 255.0) as u8,
173 (color.blue * 255.0) as u8,
174 ]);
175
176 image.put_pixel(x, y, rgb_color);
177 }
178
179 image.save(filename)
180}