SSBOのデータを抜き取る (compute shader)

音楽を作る為にSSBOのデータを抜き取ってみます。

samples = (c_float * SND_NUMSAMPLESC)()
glGetBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, sizeof(samples), byref(samples))    

この部分です。今回は音楽ですので簡単ですが、numpyとかで抜き取ってデータ加工とかにも流用できるのでは、ないのでしょうか。

from OpenGL.GL import *
from OpenGL.WGL import *
from ctypes import *
from ctypes.wintypes import *
import sys
import math

csh = """
#version 430

layout(binding=7) buffer Buf{ vec2 p[];};
layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;

#define sampleRate 44100

vec2 mainSound(float time)
{
  return vec2(sin(6.2831*440.0*fract(time)));
}

void main(){
  uint id = gl_GlobalInvocationID.x;
  float time = float(id) / sampleRate;
  p[id] =  mainSound(time);
}
"""

winmm = windll.winmm 
kernel32 = windll.kernel32
user32 = windll.user32

WS_OVERLAPPEDWINDOW = 0xcf0000
WS_VISIBLE = 0x10000000
PM_REMOVE = 1
WM_NCLBUTTONDOWN = 161
HTCLOSE = 20
VK_ESCAPE = 27
PFD_SUPPORT_OPENGL = 32
PFD_DOUBLEBUFFER = 1
WAVE_FORMAT_IEEE_FLOAT = 3
WAVE_MAPPER = -1 
TIME_SAMPLES = 2

class WAVEFORMATEX(Structure): 
    _fields_ = [ 
            ("wFormatTag", WORD), 
            ("nChannels", WORD), 
            ("nSamplesPerSec", DWORD), 
            ("nAvgBytesPerSec", DWORD), 
            ("nBlockAlign", WORD), 
            ("wBitsPerSample", WORD), 
            ("cbSize", WORD) 
    ]

class WAVEHDR(Structure): 
    pass 
LPWAVEHDR = POINTER(WAVEHDR) 
WAVEHDR._fields_ = [ 
    ("lpData", LPSTR), 
    ("dwBufferLength", DWORD), 
    ("dwBytesRecorded", DWORD), 
    ("dwUser", POINTER(DWORD)), 
    ("dwFlags", DWORD), 
    ("dwLoops", DWORD), 
    ("lpNext", LPWAVEHDR), 
    ("reserved", POINTER(DWORD)) 
]

class SMPTE(Structure):
     _fields_ = [ 
        ("hour", BYTE),
        ("min", BYTE),
        ("sec", BYTE),
        ("frame", BYTE),
        ("fps", BYTE),
        ("dummy", BYTE),
        ("pad"[2], BYTE)
     ]

class MIDI(Structure):
     _fields_ = [ 
        ("songptrpos", DWORD)
     ]

class U(Union):
     _fields_ = [ 
        ("ms", DWORD),
        ("sample", DWORD),
        ("cb", DWORD),
        ("ticks", DWORD),
        ("smpte", SMPTE),
        ("midi", MIDI),
    ]

class MMTIME(Structure):
     _fields_ = [ 
        ("wType", UINT),
        ("u", U),
    ]


duration = 10

hWnd = windll.user32.CreateWindowExA(0,0xC018,0,0,0,0,0,0,0,0,0,0)
hdc = windll.user32.GetDC(hWnd)
    
pfd = PIXELFORMATDESCRIPTOR(0,1,PFD_SUPPORT_OPENGL|PFD_DOUBLEBUFFER,32,0,0,0,0,0,0,0,0,0,0,0,0,0,32,0,0,0,0,0,0,0)
SetPixelFormat(hdc, ChoosePixelFormat(hdc, pfd), pfd)
hGLrc = wglCreateContext(hdc)
wglMakeCurrent(hdc, hGLrc)

SAMPLE_RATE     = 44100
SND_NUMCHANNELS = 2
SND_NUMSAMPLES  = math.ceil(duration*SAMPLE_RATE)
SND_NUMSAMPLESC = SND_NUMSAMPLES*SND_NUMCHANNELS
samples = (c_float * SND_NUMSAMPLESC)()

computeProg = glCreateProgram()
shader = glCreateShader(GL_COMPUTE_SHADER)
glShaderSource(shader, csh)
glCompileShader(shader)
if glGetShaderiv(shader, GL_COMPILE_STATUS) != GL_TRUE:
    raise RuntimeError(glGetShaderInfoLog(shader).decode())
glAttachShader(computeProg, shader)
glLinkProgram(computeProg)

vbo = glGenBuffers(1)
glBindBuffer(GL_SHADER_STORAGE_BUFFER, vbo)
glBufferData(GL_SHADER_STORAGE_BUFFER, sizeof(samples), None, GL_STATIC_DRAW)
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 7, vbo)
    
glUseProgram(computeProg)
glDispatchCompute(SND_NUMSAMPLESC // 128, 1, 1)
    
glGetBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, sizeof(samples), byref(samples))    
    
wh = WAVEHDR(cast(samples, LPSTR), sizeof(samples)) 
wfx = WAVEFORMATEX(
    WAVE_FORMAT_IEEE_FLOAT,
    SND_NUMCHANNELS,
    SAMPLE_RATE,
    SAMPLE_RATE*sizeof(c_float)*SND_NUMCHANNELS,
    sizeof(c_float)*SND_NUMCHANNELS,
    sizeof(c_float)*8,
    0
)
wh = WAVEHDR(cast(samples, LPSTR), sizeof(samples)) 
hwo = HANDLE(0) 
winmm.waveOutOpen(byref(hwo), WAVE_MAPPER, byref(wfx), 0, 0, 0)
winmm.waveOutPrepareHeader(hwo, byref(wh), sizeof(WAVEHDR))
winmm.waveOutWrite(hwo, byref(wh), sizeof(WAVEHDR))

mmt = MMTIME(TIME_SAMPLES)
msg = MSG()
lpmsg = pointer(msg)
done = False
while done==False:
    while user32.PeekMessageA(lpmsg, 0, 0, 0, PM_REMOVE):
        if (msg.message == WM_NCLBUTTONDOWN and msg.wParam == HTCLOSE): done = True
        user32.DispatchMessageA(lpmsg)
    if(user32.GetAsyncKeyState(VK_ESCAPE)):  done = True
    winmm.waveOutGetPosition(hwo, byref(mmt), sizeof(mmt))
    sys.stdout.write("\r TIME : %f" %(mmt.u.sample / SAMPLE_RATE,))
    sys.stdout.flush()
    if (mmt.u.sample == SND_NUMSAMPLES): done = True
    
winmm.waveOutClose(hwo)
wglMakeCurrent(0, 0)
wglDeleteContext(hGLrc)
user32.ReleaseDC(hWnd, hdc)
user32.PostQuitMessage(0)
user32.DestroyWindow(hWnd)