Source

raytracer / rays.ml

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
(* ray.ml: a simple raytracer *)

include Json

type vector3 = float * float * float
type color = int * int * int
type normal = Normal of vector3
type surface = {surface_color:color;reflectance:float;albedo:float}
type light = {light_center:vector3; light_color:color}
type sphere = {center:vector3; radius:float; sph_surf:surface; sph_id:int}
type plane = {plane_normal:normal; offset:float; plane_surf:surface;
	      plane_id:int}
type shape = Sphere of sphere | Plane of plane
type intersection = {location:vector3;normal:normal;obj:shape}
type camera = {origin:vector3; out:normal; up:normal; width:int;
	       height:int; xangle:float; yangle:float}

(* color functions and constants *)

let rgb_black : color = (0,0,0)
let rgb_white : color = (255,255,255)
let rgb_red : color = (255,0,0)
let rgb_green : color = (0,255,0)
let rgb_blue : color = (0,0,255)
let rgb_yellow : color = (255,255,0)
let rgb_cyan : color = (0,255,255)

let rgb_r ((r,g,b):color) = r
let rgb_g ((r,g,b):color) = g
let rgb_b ((r,g,b):color) = b
let rgb_coeff ((r,g,b):color) (coeff:float) = 
  let mult i = int_of_float ((float_of_int i) *. coeff) in
    (mult r, mult g, mult b)
let rgb_mix ((r1,g1,b1):color) ((r2,g2,b2):color) : color =
  let mix a b = (a*b)/255 in
    (mix r1 r2, mix g1 g2, mix b1 b2)
let rgb_sum ((r1,g1,b1):color) ((r2,g2,b2):color) : color =
  (r1+r2,g1+g2,b1+b2)
let rgb_arith_mix (c1:color) (c2:color) (mix:float) : color =
  assert (mix <= 1.0 && mix >= 0.0);
  rgb_sum (rgb_coeff c1 mix) (rgb_coeff c2 (1.0 -. mix))
let rgb_chop ((r,g,b):color) : color =
  let chop i = min 255 (max 0 i) in
  (chop r, chop g, chop b)

let fabs (f:float) : float = if f > 0.0 then f else 0.0 -. f
let sq x = x *. x 

(* some convenient surfaces *)

let surf_shiny_green = {surface_color=rgb_green; reflectance=0.0; albedo=1.0}
let surf_shiny_red = {surface_color=rgb_red; reflectance=0.5; albedo=0.5}
let surf_shiny_blue = {surface_color=rgb_blue; reflectance=0.1; albedo=0.5}
let surf_shiny_yellow = {surface_color=rgb_yellow;reflectance=0.9; albedo=0.5}
let surf_shiny_cyan = {surface_color=rgb_cyan;reflectance=0.9; albedo=0.5}

(* vector3 functions *)

let v3_origin = (0.0,0.0,0.0)
let v3_l2distsq ((x1,y1,z1):vector3) ((x2,y2,z2):vector3) =
  let xdiff = x1 -. x2 in
  let ydiff = y1 -. y2 in
  let zdiff = z1 -. z2 in
    (sq xdiff) +. (sq ydiff) +. (sq zdiff)
let v3_l2dist p1 p2 = sqrt (v3_l2distsq p1 p2)
let v3_magnitudesq (v3:vector3) : float = v3_l2distsq v3 v3_origin
let v3_magnitude (v3:vector3) : float = sqrt (v3_magnitudesq v3)
let v3_mul ((x,y,z):vector3) (f:float) : vector3 = (x*.f, y*.f, z*.f)
let v3_div (v3:vector3) (f:float) : vector3 =
  let recip = 1.0 /. f in
    v3_mul v3 recip
let v3_neg ((x,y,z):vector3) : vector3 = (-.x, -.y, -.z)
let v3_add ((x1,y1,z1):vector3) ((x2,y2,z2):vector3) : vector3 =
  (x1+.x2, y1+. y2, z1 +. z2)
let v3_sub (v3a:vector3) (v3b:vector3) : vector3 = v3_add v3a (v3_neg v3b)
let v3_dot ((x1,y1,z1):vector3) ((x2,y2,z2):vector3) : float =
  (x1*.x2) +. (y1*.y2) +. (z1*.z2)
let v3_cross ((a1,a2,a3):vector3) ((b1,b2,b3):vector3) : vector3 =
  (a2*.b3 -. a3*.b2, a3*.b1 -. a1*.b3, a1*.b2 -. a2*.b1)  
let v3_norm (v3:vector3) : normal = Normal (v3_div v3 (v3_magnitude v3))

(* shape intersection handlers *)

let compare_shapes (s1:shape) (s2:shape) : bool =
  match (s1,s2) with
    | (Sphere sph1, Sphere sph2) -> sph1.sph_id = sph2.sph_id
    | (Plane p1, Plane p2) -> p1.plane_id = p2.plane_id
    | _ -> false

let shape_surface (s:shape) : surface =
  match s with
      Sphere sph -> sph.sph_surf
    | Plane p -> p.plane_surf


let sphere_intersect (sph:sphere) (src:vector3) (dir:normal) :
    intersection option =
  let center = v3_sub sph.center src in
  let l = match dir with Normal n -> n in
  let lcdot = v3_dot l center in
  let lcdot2 = sq lcdot in
  let magc2 = v3_magnitudesq center in
  let r2 = sq sph.radius in
  let radical = lcdot2 -. magc2 +. r2 in
  let solution () =
    let sqrtrad = sqrt radical in
    let plusmag = lcdot +. sqrtrad in
    let minusmag = lcdot -. sqrtrad in
    let plus = v3_add src (v3_mul l plusmag) in
    let minus = v3_add src (v3_mul l minusmag) in
    let plus_closer = (v3_l2distsq src plus) < (v3_l2distsq src minus) in
    let loc = if plus_closer then plus else minus in
    let norm = v3_norm (v3_sub loc sph.center) in
      Some {location=loc;normal=norm;obj=(Sphere sph)}
  in
    if radical >= 0.0 then solution() else None

let plane_intersect (p:plane) (src:vector3) (dir:normal) : intersection option =
  let l = match dir with Normal n -> n in
  let n = match p.plane_normal with Normal n -> n in
  let denom = v3_dot l n in
    if (fabs denom) < 0.000000001 then None else (
      let t = ((-.p.offset) -. (v3_dot src n)) /. denom in
      let loc = v3_add (v3_mul l (t)) src in
      let norm = p.plane_normal in
	if t > 0.0 then
	  (Some {location=loc; normal=norm; obj=Plane p})
	else 
	  None)

let shape_intersect (src:vector3) (dir:normal) (shape:shape) :
    intersection option =
  match shape with
    | Sphere s -> sphere_intersect s src dir
    | Plane p -> plane_intersect p src dir

let closest_intersection_excluded (source_isect:intersection option) 
    (shapes:shape list) (src:vector3) (dir:normal) : intersection option =
  let isects =
    let filter_isect acc s =
      match shape_intersect src dir s with
	  None -> acc
	| Some isect -> (
	    match source_isect with
	      | None -> isect::acc
	      | Some orig_isect -> (
		  if (compare_shapes orig_isect.obj s) &&
		    ((v3_l2distsq orig_isect.location isect.location)
		     < 0.0000001) then acc else isect::acc))
    in
      List.fold_left filter_isect [] shapes
  in
  let choose_closest (prev:intersection*float)
			(next:intersection) =
    let dist = v3_l2distsq src next.location in
    let (best,best_dist) = prev in
      if dist < best_dist then (next,dist) else prev
  in
  let extract_isect (best,best_dist) = best in
    match isects with
      | [] -> None
      | _ -> Some (
	  let hd = List.hd isects in
	  let hd_dist = v3_l2distsq hd.location src in
	  let best_pair = 
	    List.fold_left choose_closest (hd,hd_dist) (List.tl isects)
	  in
	    extract_isect best_pair)

let closest_intersection = closest_intersection_excluded None

(* camera functions *)

let build_rays (cam:camera) : normal list =
  let d_theta_x = cam.xangle /. (float_of_int cam.width) in
  let d_theta_y = cam.yangle /. (float_of_int cam.height) in
  let halfx = cam.xangle /. 2.0 in
  let halfy = cam.yangle /. 2.0 in
  let extract_norm n = match n with Normal v -> v in
  let out = extract_norm cam.out in
  let up = extract_norm cam.up in
  let right = v3_cross out up in
  let rec loop ret x y = 
    match x with
      | 0 -> (
	  match y with
	    | 0 -> ret
	    | _ -> loop ret cam.width (y-1))
      | _ -> (
	  let xtheta = (float_of_int x)*.d_theta_x -. halfx in
	  let ytheta = (float_of_int y)*.d_theta_y -. halfy in
	  let right_comp = v3_mul right (sin xtheta) in
	  let up_comp = v3_mul up (sin ytheta) in
	  let next = v3_norm (v3_add (v3_add out up_comp) right_comp) in
	    loop (next::ret) (x-1) y)
  in
    loop [] cam.width cam.height

let illuminated (isect:intersection) (shapes:shape list) (light:light) : bool =
  let dir = v3_norm (v3_sub light.light_center isect.location) in
  let closest = 
    closest_intersection_excluded (Some isect) shapes isect.location dir
  in
    match closest with
	None -> true
      | Some isect2 -> (
	  let isectd = v3_l2dist isect2.location isect.location in
	  let lightd = v3_l2dist light.light_center isect.location in
	    isectd > lightd)

let handle_ambient (isect:intersection) (ambience:float) : color =
  rgb_coeff (shape_surface isect.obj).surface_color ambience

let rec render_ray (depth:int) (max_depth:int) (src:vector3)
    (shapes:shape list) (lights:light list) (bg:color) (ambience:float)
    (ray:normal) : color =
  let handle_illuminated (isect:intersection) (ray:normal) (light:light)
      : color =
    let strip_norm = function Normal n -> n in
    let reflect_dir v3 norm =
      let normal_comp = v3_dot v3 norm in
	v3_add v3 (v3_mul norm (-.2.0*.normal_comp))
    in
    let light_dir = v3_norm (v3_sub light.light_center isect.location) in
    let light_norm = strip_norm light_dir in
    let isect_norm = strip_norm isect.normal in
    let ray_norm = strip_norm ray in
    let lambertian_coeff = max (v3_dot light_norm isect_norm) 0.0 in
    let light_in = rgb_coeff light.light_color lambertian_coeff in
    let surf = shape_surface isect.obj in
    let lambert_color = rgb_mix light_in surf.surface_color in
    let light_reflect_dir = reflect_dir light_norm isect_norm in
    let specular_coeff = v3_dot light_reflect_dir ray_norm in
    let specular_color = rgb_coeff light.light_color specular_coeff in
    let own_color = 
      rgb_chop (rgb_sum lambert_color (rgb_coeff specular_color surf.albedo))
    in
    let ray_reflect_dir = v3_norm (reflect_dir ray_norm isect_norm) in
    let reflect_color =
      if surf.reflectance > 0.01 && depth < max_depth then
	let rrd = strip_norm ray_reflect_dir in
	let next_src = v3_add isect.location (v3_mul rrd 0.000001) in
	  render_ray (depth+1) max_depth next_src shapes
	    lights bg ambience ray_reflect_dir
      else rgb_black
    in
      if depth < max_depth then
	rgb_arith_mix reflect_color own_color surf.reflectance
      else own_color
  in
  let closest = closest_intersection shapes src ray in
    match closest with
	None -> bg
      | Some isect -> (
	  let handle_light color light =
	    let is_illum = illuminated isect shapes light in
	      if is_illum then (
		rgb_sum color (handle_illuminated isect ray light))
	      else color
	  in
	    rgb_sum (handle_ambient isect ambience)
	      (List.fold_left handle_light rgb_black lights))
	  
let render_scene (camera:camera) (shapes:shape list) (lights:light list)
    (bg:color) (ambience:float) (max_depth:int) : color list =
  let rays = build_rays camera in
  let f = render_ray 0 max_depth camera.origin shapes lights bg ambience in
  let rec loop ret rem =
    match rem with
      | [] -> List.rev ret
      | hd::rest -> loop ((f hd)::ret) rest
  in
    loop [] rays

let dump_bitmap (out:out_channel) (camera:camera) (colors:color list)  : unit =
  let n_pixels = List.length colors in
  let image_size = n_pixels * 3 in
  let header_size = 54 in
  let total_size = image_size + header_size in
  let out_byte (i:int) = output_byte out i in
  let out_int (i:int) =
    out_byte i;
    out_byte (i asr 8);
    out_byte (i asr 16);
    out_byte (i asr 24)
  in
  let out_short (i:int) =
    out_byte i;
    out_byte (i asr 8);
  in    
  let out_color ((r,g,b):color) =
    let out_c (i:int) = out_byte (min 255 (max 0 i)) in
    out_c b;
    out_c g;
    out_c r;
  in
  let write_header () =
    output_char out 'B';
    output_char out 'M';
    out_int total_size;
    out_int 0; (* used by BMP creating application *)
    out_int 54; (* offset to bitmap data *)
    out_int 40; (* DIB header size -- using Windows V3 version *)
    out_int camera.width;
    out_int camera.height;
    out_short 1; (* number of color planes *)
    out_short 24; (* color bits *)
    out_int 0; (* compression mode - no compression *)
    out_int image_size;
    out_int (1 lsl 24); (* horizontal resolution, arbitrarily chosen *)
    out_int (1 lsl 24); (* vertical resolution, arbitrarily chosen *)
    out_int 0; (* default number of colors in palette *)
    out_int 0; (* number of important colors, 0 = all colors important *)
  in
  let write_data () =
    assert (camera.width mod 4 = 0);
    assert (camera.height mod 4 = 0);
    List.iter out_color colors
  in
    write_header();
    write_data()

let write_bitmap (filename:string) (camera:camera) (colors:color list) : unit =
  let out = open_out_bin filename in
    dump_bitmap out camera colors;
    close_out out

(* parsing objects from JSON *)

let extract_triple (f: Json.json -> 'a) : Json.json -> 'a * 'a * 'a =
  let lookup (js:Json.json) =
    let lu i = f (Json.lookup (Json.Int i) js) in
      (lu 0, lu 1, lu 2)
  in
    lookup

let extract_color = extract_triple Json.force_int

let extract_v3 = extract_triple Json.force_float

let str_lookup (js:Json.json) (s:string) : Json.json =
  Json.lookup (Json.String s) js

let parse_surface_js (js:Json.json) : surface = 
  let lu = str_lookup js in
  let surface_color = extract_color (lu "surface_color") in
  let reflectance = Json.force_float (lu "reflectance") in
  let albedo = Json.force_float (lu "albedo") in
    {surface_color=surface_color; reflectance=reflectance; albedo=albedo}

let parse_sphere_js (js:Json.json) : shape =
  let lu = str_lookup js in
  let center = extract_v3 (lu "center") in
  let radius = Json.force_float (lu "radius") in
  let id = Json.force_int (lu "sph_id") in
  let surface = parse_surface_js (lu "sph_surf") in
    Sphere {center=center; radius=radius; sph_surf=surface; sph_id=id}

let parse_plane_js (js:Json.json) : shape = 
  let lu = str_lookup js in
  let plane_normal = v3_norm (extract_v3 (lu "plane_normal")) in
  let offset = Json.force_float (lu "offset") in
  let plane_surf = parse_surface_js (lu "plane_surf") in
  let plane_id = Json.force_int (lu "plane_id") in
    Plane {plane_normal=plane_normal; offset=offset; plane_surf=plane_surf;
	   plane_id=plane_id}

exception Unknown_shape of string

let parse_shape_js (js:Json.json) : shape =
  let lu = str_lookup js in 
  let name = Json.force_string (lu "name") in
  let fields = lu "fields" in
    match name with
      | "sphere" -> parse_sphere_js fields
      | "plane" -> parse_plane_js fields
      | _ -> raise (Unknown_shape name)

let parse_light_js (js:Json.json) : light =
  let lu = str_lookup js in
  let light_center = extract_v3 (lu "light_center") in
  let light_color = extract_color (lu "light_color") in
    {light_center=light_center; light_color=light_color}

let parse_camera_js (js:Json.json) : camera =
  let lu = str_lookup js in
  let origin = extract_v3 (lu "origin") in
  let out = v3_norm (extract_v3 (lu "out")) in
  let up = v3_norm (extract_v3 (lu "up")) in
  let width = Json.force_int (lu "width") in
  let height = Json.force_int (lu "height") in
  let xangle = Json.force_float (lu "xangle") in
  let yangle = Json.force_float (lu "yangle") in
    {origin=origin; out=out; up=up; width=width; height=height;
     xangle=xangle; yangle=yangle}

let parse_scene (js:Json.json) : (camera * (shape list) * (light list)
				  * color * float * int) =
  let lu = str_lookup js in
  let lights = List.map parse_light_js (Json.force_array (lu "lights")) in
  let shapes = List.map parse_shape_js (Json.force_array (lu "shapes")) in
  let camera = parse_camera_js (lu "camera") in
  let bg = extract_color (lu "bg") in
  let ambience = Json.force_float (lu "ambience") in
  let bounces = Json.force_int (lu "bounces") in
    (camera, shapes, lights, bg, ambience, bounces)
  

let shapes = [Sphere {center=(10.0,0.0,0.0); radius=2.0;
		      sph_surf=surf_shiny_green; sph_id=0};
	      Plane {plane_normal=v3_norm (0.0,1.0,0.0); offset=(5.0);
		    plane_surf=surf_shiny_blue; plane_id=1};
	      Plane {plane_normal=v3_norm (-.1.0,0.0,0.5); offset=(15.0);
		     plane_surf=surf_shiny_cyan; plane_id=2};
	      Plane {plane_normal=v3_norm (0.0,0.0,-1.0); offset=10.0;
		     plane_surf=surf_shiny_yellow; plane_id=3}]
let lights = [{light_center=(0.0,10.0,0.0); light_color=(200,200,200)};
	      {light_center=(3.0,0.0,-.4.0);light_color=(255,100,100)}]
let camera = {origin=v3_origin; out=v3_norm (1.0,0.0,0.0);
	      up=v3_norm (0.0,1.0,0.0); width=1024; height=768;
	      xangle=2.0; yangle=1.5}

(*let colors = render_scene camera shapes lights rgb_red 0.07 1
let out = write_bitmap "rendered.bmp" camera colors*)

let load_and_render_scene (fname:string) (out:string) : unit =
  let npt = open_in fname in
  let buflen = (in_channel_length npt) in
  let buf = String.make buflen ' ' in
  let n_read = input npt buf 0 buflen in
    assert (n_read = buflen);
  let js = Json.loads buf in
  let (camera,shapes,lights,bg,ambience,bounces) = parse_scene js in
  let colors = render_scene camera shapes lights bg ambience bounces in
    write_bitmap out camera colors

let dtest =
  assert (fabs ((v3_l2dist (3.0, 0.0, 0.0) (0.0, 4.0, 0.0)) -. 5.0) < 0.00001)

let main = 
  if Array.length Sys.argv != 3 then (
    print_endline "wrong number of arguments")
  else (
    let fname = Sys.argv.(1) in
    let out = Sys.argv.(2) in
      load_and_render_scene fname out)