diff --git a/shared/src/core/image.rs b/shared/src/core/image.rs index fe6d702..a9025ee 100644 --- a/shared/src/core/image.rs +++ b/shared/src/core/image.rs @@ -69,10 +69,55 @@ impl PixelFormat { #[repr(C)] #[derive(Clone, Copy, Debug)] -pub enum Pixels { - U8(Ptr), - F16(Ptr), - F32(Ptr), +pub struct Pixels { + ptr: Ptr, + format: PixelFormat, +} + +impl Pixels { + pub fn new_u8(ptr: Ptr) -> Self { + Self { + ptr, + format: PixelFormat::U8, + } + } + + pub fn new_f16(ptr: Ptr) -> Self { + Self { + ptr: Ptr::from_raw(ptr.as_raw() as *const u8), + format: PixelFormat::F16, + } + } + + pub fn new_f32(ptr: Ptr) -> Self { + Self { + ptr: Ptr::from_raw(ptr.as_raw() as *const u8), + format: PixelFormat::F32, + } + } + + unsafe fn read_u8(&self, byte_offset: usize) -> u8 { + unsafe { *self.ptr.as_raw().add(byte_offset) } + } + + unsafe fn read_f16(&self, elem_offset: usize) -> u16 { + let byte_offset = elem_offset * 2; + unsafe { *(self.ptr.as_raw().add(byte_offset) as *const u16) } + } + + unsafe fn read_f32(&self, elem_offset: usize) -> f32 { + let byte_offset = elem_offset * 4; + unsafe { *(self.ptr.as_raw().add(byte_offset) as *const f32) } + } + + // Unsure if ill need this + // pub unsafe fn read(&self, offset: usize) -> Float { + // match self.format { + // PixelFormat::U8 => unsafe { self.read_u8(offset) as Float / 255.0 }, + // PixelFormat::F16 => unsafe { f16_to_f32_software(self.read_f16(offset)) }, + // PixelFormat::F32 => unsafe { self.read_f32(offset) }, + // } + // } } #[repr(C)] @@ -170,18 +215,18 @@ impl ImageAccess for DeviceImage { return 0.; } - let offset = self.pixel_offset(p) + c as u32; + let offset = (self.pixel_offset(p) + c as u32) as usize; unsafe { - match self.pixels { - Pixels::U8(ptr) => { - let raw_val = *ptr.add(offset as usize); - self.base().encoding.to_linear_scalar(raw_val) + match self.pixels.format { + PixelFormat::U8 => { + let raw_val = self.pixels.read_u8(offset); + self.base.encoding.to_linear_scalar(raw_val) } - Pixels::F16(ptr) => { - let raw_val = *ptr.add(offset as usize); + PixelFormat::F16 => { + let raw_val = self.pixels.read_f16(offset); f16_to_f32_software(raw_val) } - Pixels::F32(ptr) => *ptr.add(offset as usize), + PixelFormat::F32 => self.pixels.read_f32(offset), } } }