photon/render_image_example.py at main · bvillasen/photon · GitHub

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
import sys, time, os
from os import listdir
from os.path import isfile, join
import h5py as h5
import numpy as np
import pycuda.gpuarray as gpuarray
import matplotlib.colors as cl
import matplotlib.cm as cm
from PIL import Image

#Add Modules from other directories
currentDirectory = os.getcwd()
srcDirectory = currentDirectory + "/src/"
dataDirectory = currentDirectory + "/data_src/"
sys.path.extend([ srcDirectory, dataDirectory ] )
import volumeRender_image as volumeRender
from cudaTools import setCudaDevice, getFreeMemory, gpuArray3DtocudaArray, np3DtoCudaArray
from data_functions import *
import gpu_data
from tools import create_directory

from IPython import embed

#Select CUDA Device
useDevice = 0


#data_dir = '/home/bruno/Desktop/ssd_0/data/'
data_dir = '/home/xavier/Desktop/mhws/'
# data_dir = '/raid/bruno/data/'
# data_dir = '/data/groups/comp-astro/bruno/'
output_dir = os.path.join(data_dir, 'render_images')
#output_dir = '/home/bruno/Desktop/'
create_directory( output_dir )

# Data from file
input_file_name = '/home/xavier/Desktop/mhws/mhws_for_bruno.h5'

file = h5.File( input_file_name, 'r' )
grid = file['grid'][...]
#embed(header='41 of image example')

n_cut = 70
grid = grid[n_cut:-n_cut,::]

n_fields_to_render = 1
data_parameters = {
'type': 'field',
'data': grid,
#  'dims':[512, 512, 512],
  'log_data': False,
  'normalization':'local',
  'min_data': 0,
  'n_border':1  # was 0
}

n_image = 0
data_to_render_list = [ get_Data_to_Render(
    data_parameters ) for i in range(n_fields_to_render)]

volumeRender.render_parameters[0] = { 'transp_type':'flat',
                                     'density':.05,
                                     "brightness":1.2,
                                     'transfer_offset': 0.0,
                                     'transfer_scale': 1. }
volumeRender.render_parameters[0]['colormap'] = 'jet'
volumeRender.render_parameters[0]['transp_type'] = 'linear'
volumeRender.render_parameters[0]['transp_min'] = 0.0
volumeRender.render_parameters[0]['transp_max'] = 1.0
volumeRender.render_parameters[0]['output_transfer'] = output_dir

#Get Dimensions of the data to render
nz, ny, nx = data_to_render_list[0].shape

#Size of the Image
image_height, image_width = 1024, 2048

#Initialize the volumeRender
volumeRender.image_width = image_width
volumeRender.image_height = image_height
volumeRender.nTextures = n_fields_to_render
volumeRender.nWidth = nx
volumeRender.nHeight = ny
volumeRender.nDepth = nz

#initialize pyCUDA context
cudaDevice = setCudaDevice( devN=useDevice, usingAnimation=False )

# Initialize the Volume Render Functions
volumeRender.initCUDA()

# Initialize GPU Data Functions
gpu_data.Init_GPU_Data()

#Initialize all gpu data
gpu_array_list, gpu_array_fixed_list,  copyToScreen_list = gpu_data.Initialize_GPU_Data(  data_to_render_list, volumeRender )

# Copy the data to the gpu
for i in range(n_fields_to_render):  copyToScreen_list[i]

# Create Array for the Image
image_data_h = np.zeros( [4, image_height, image_width],
                        dtype=np.int32 )
image_data_d = gpuarray.to_gpu( image_data_h )

# Apply a rotation of the data
volumeRender.viewRotation = [-11.6, -44.4, 0. ]

# Render the Image
volumeRender.render_image( render_to=image_data_d )
image_data_h = image_data_d.get()

# Reshape the image to RGBA
rgba_data = np.zeros( [image_height, image_width, 4],
                     dtype=np.uint8 )
for i in range(4): rgba_data[:,:, i] = image_data_h[i,:,:]

# Save Image to PNG
out_file_name = os.path.join(output_dir, f'render_image.png')
image_rgba = Image.fromarray( rgba_data )
image_rgb = image_rgba.convert('RGB')
image_rgb.save( out_file_name )
print(f'Saved Image: {out_file_name}' )