Try to offload tight loop to rust (for now it segfaults)

[?]

Jul 4, 2021, 6:47 PM

6AXPZL5PQMW5P6BPRTGWO3AOWAU5OC33UYBY3TVOHUVQXCRVVTPQC

Dependencies

[2] IGYI5RVV Figure out how to draw on the sdl2 window
[3] VKA5CCGC Add pysdl2 to the trainmodel dependencies so I'll be able to visualize what's going on.
[4] ROQCAPZJ Begin function for showing the map (for now just opens SDL window)
[5] ZRPV3GAJ Model trainer: Add function to make tensor of cartesian coordinates for drawing the map (equirectangular projection)
[6] QBDHX7BH Add a makefile to express python dependencies since pipenv doesn't like tensorflow
[*] X5U7KRUI Parse some of the data into tensors

Change contents

replacement in trainmodel/Makefile at line 1

[3.11]→[3.12:48](∅→∅)

venv:
virtualenv venv

setup: venv

[3.11]

[3.48]

setup: venv src/helpers.so
edit in trainmodel/Makefile at line 6

[3.54]

venv:
virtualenv venv

src/helpers.so: src/helpers.rs
rustc --crate-name helpers --edition=2018 src/helpers.rs -o src/helpers.so --crate-type cdylib --emit=link -C embed-bitcode=no -C link-args="$$(sdl2-config --libs)"
edit in trainmodel/src/drawmap.py at line 2

[2.14]→[2.14:15](∅→∅)
edit in trainmodel/src/drawmap.py at line 3

[3.12]

[3.14]

import os
edit in trainmodel/src/drawmap.py at line 18

[2.168]

[2.168]

print("rendered")
print(image)
print(image.contents)
print("pixels: " + hex(image.contents.pixels))
print(ctypes.cast(image.contents.pixels, ctypes.POINTER(ctypes.c_uint32)).contents)
edit in trainmodel/src/drawmap.py at line 24

[2.225]

[2.225]

print("blitted")
edit in trainmodel/src/drawmap.py at line 48

[3.407]

[2.277]

helpers = ctypes.cdll.LoadLibrary(os.path.dirname(os.path.realpath(__file__)) + '/helpers.so')
helpers.render_tensor.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int]
helpers.render_tensor.restype = ctypes.POINTER(sdl2.surface.SDL_Surface)
edit in trainmodel/src/drawmap.py at line 52

[2.278]→[2.278:460](∅→∅)

# TODO: refactor to only do a few pixels at a time because python is really slow
# Keeping this function here in the meantime because it wasn't easy to figure
# out how to write it.
replacement in trainmodel/src/drawmap.py at line 54

[2.516]→[2.516:976](∅→∅)

height = shape[0]
width = shape[1]
surface = sdl2.surface.SDL_CreateRGBSurface(0, width, height, 32, 0, 0, 0, 0)
pixels = sdl2.ext.pixels2d(surface.contents, False)
# pixels are 4 bytes: b g r 0
for y in range(0,height):
row = source[y]
d_row = pixels[y]
for x in range(0,width):
pixel = row[x]
d_row[x] = pixel[0] * 65536 + pixel[1] * 256 + pixel[2]
print(y)
return surface

[2.516]

return helpers.render_tensor(ctypes.c_void_p(tf.cast(source, tf.uint8).numpy().ctypes.data), ctypes.c_int(shape[1]), ctypes.c_int(shape[0]), ctypes.c_int(shape[2]))
file addition: helpers.rs (----------)

[8.6]

use std::ffi::c_void;

#[repr(C)]
pub struct SDLPixelFormat {
data: c_void,
}

#[repr(C)]
pub struct SDLRect {
x: i16,
y: i16,
w: u16,
h: u16,
}

#[repr(C)]
pub struct SDLSurface {
flags: u32,
format: *const SDLPixelFormat,
w: i32,
h: i32,
pitch: i32,
pixels: *mut c_void,
userdata: *mut c_void,
locked: i32,
lock_data: *mut c_void,
clip_rect: SDLRect,
map: *mut c_void,

refcount: i32,
}

#[link(name = "SDL")]
extern{
#[link_name = "SDL_MapRGBA"]
fn c_sdl_map_rgba(format: *const SDLPixelFormat, r: u8, g: u8, b: u8, a: u8) -> u32;
#[link_name = "SDL_CreateRGBSurface"]
fn c_sdl_create_rgb_surface(flags: u32, width: u32, height: u32, depth: u32, r_mask: u32, g_mask: u32, b_mask: u32, a_mask: u32) -> *mut SDLSurface;
}

fn map_rgba(format: &SDLPixelFormat, r: u8, g: u8, b: u8, a: u8) -> u32 {
unsafe {
c_sdl_map_rgba(format as *const SDLPixelFormat, r, g, b, a)
}
}

#[no_mangle]
pub extern "C" fn render_tensor(data: *const u8, width: u32, height: u32, channels: i32) -> *mut SDLSurface {
let surface = unsafe {
c_sdl_create_rgb_surface(0, width, height, 32, 0xff, 0xff00, 0xff0000, 0xff000000)
};
println!("surface: {:?}, pixels: {:?}, w: {}, h: {}", surface, unsafe { (*surface).pixels }, unsafe { (*surface).w }, unsafe { (*surface).h });
let width = width as usize;
let height = height as usize;
let channels = channels as usize;
let pitch = unsafe {
(*surface).pitch as usize / 4
};
let data = unsafe {
std::slice::from_raw_parts(data, width * height * channels)
};
let pixels: &mut [u32] = unsafe {
std::slice::from_raw_parts_mut((*surface).pixels.cast(), height * pitch)
};
let format = unsafe {
(*surface).format.as_ref().unwrap()
};
for y in 0 .. height {
let data_rowstart = y * width * (channels as usize);
let pixels_rowstart = y * pitch;
for x in 0 .. width {
let data_pixstart = data_rowstart + x * channels;
let (red, green, blue, alpha) = match channels {
1 => {
let grey = data[data_pixstart];
(grey, grey, grey, 255)
}
3 => {
let red = data[data_pixstart];
let green = data[data_pixstart + 1];
let blue = data[data_pixstart + 2];
(red, green, blue, 255)
}
4 => {
let red = data[data_pixstart];
let green = data[data_pixstart + 1];
let blue = data[data_pixstart + 2];
let alpha = data[data_pixstart + 3];
(red, green, blue, alpha)
}
_ => panic!("We only handle 1, 3, and 4 channels!")
};
let color_value = map_rgba(format, red, green, blue, alpha);
pixels[pixels_rowstart + x] = color_value;
}
}
println!("Surface = {:?}, pixels = {:?}", surface, unsafe { (*surface).pixels });
surface
}