1- __version__ = "1 .0"
1+ __version__ = "2 .0"
22
33from re import M
44from meshroom .core import desc
55from meshroom .core .utils import VERBOSE_LEVEL
6-
7-
8- class RaftOFNodeSize (desc .MultiDynamicNodeSize ):
9- def computeSize (self , node ):
10- from pathlib import Path
11- import itertools
12-
13- input_path_param = node .attribute (self ._params [0 ])
14- extension_param = node .attribute (self ._params [1 ])
15-
16- input_path = input_path_param .value
17- extension = extension_param .value
18- include_suffixes = [extension .lower (), extension .upper ()]
19-
20- size = 1
21- if Path (input_path ).is_dir ():
22- image_paths = list (itertools .chain (* (Path (input_path ).glob (f'*.{ suffix } ' ) for suffix in include_suffixes )))
23- size = len (image_paths )
24- elif node .attribute (self ._params [0 ]).isLink :
25- size = node .attribute (self ._params [0 ]).getLinkParam ().node .size
26-
27- return size
6+ from pyalicevision import parallelization as avpar
287
298
309class RaftOFBlockSize (desc .Parallelization ):
@@ -46,25 +25,16 @@ class RaftOF(desc.Node):
4625
4726 gpu = desc .Level .INTENSIVE
4827
49- size = RaftOFNodeSize ([ ' inputImages' , 'inputExtension' ] )
28+ size = avpar . DynamicViewsSize ( " inputImages" )
5029 parallelization = RaftOFBlockSize ()
5130
5231 inputs = [
5332 desc .File (
5433 name = "inputImages" ,
5534 label = "Input Images" ,
56- description = "Input images to estimate the depth from. Folder path or sfmData filepath" ,
35+ description = "Input images to estimate the depth from, sfmData filepath" ,
5736 value = "" ,
5837 ),
59- desc .ChoiceParam (
60- name = "inputExtension" ,
61- label = "Input Extension" ,
62- description = "Extension of the input images. This will be used to determine which images are to be used if \n "
63- "a directory is provided as the input. If \" \" is selected, the provided input will be used as such." ,
64- values = ["jpg" , "jpeg" , "png" , "exr" ],
65- value = "exr" ,
66- exclusive = True ,
67- ),
6838 desc .File (
6939 name = "modelPath" ,
7040 label = "Model Path" ,
@@ -109,20 +79,18 @@ class RaftOF(desc.Node):
10979 value = "{nodeCacheFolder}" ,
11080 ),
11181 desc .File (
112- name = "OpticalFlow " ,
113- label = "Optical Flow" ,
114- description = "Output optical flow images" ,
82+ name = "OpticalFlowUp " ,
83+ label = "Optical Flow Up " ,
84+ description = "Output up optical flow images" ,
11585 semantic = "image" ,
11686 value = lambda attr : "{nodeCacheFolder}/<FILESTEM>.exr" ,
117- group = "" ,
118- )
87+ ),
11988 ]
12089
12190 def preprocess (self , node ):
122- extension = node .inputExtension .value
12391 input_path = node .inputImages .value
12492
125- image_paths = get_image_paths_list (input_path , extension )
93+ image_paths = get_image_paths_list (input_path )
12694
12795 if len (image_paths ) == 0 :
12896 raise FileNotFoundError (f'No image files found in { input_path } ' )
@@ -167,11 +135,12 @@ def processChunk(self, chunk):
167135 chunk .logger .info (f'Starting computation on chunk { chunk .range .iteration + 1 } /{ chunk .range .fullSize // chunk .range .blockSize + int (chunk .range .fullSize != chunk .range .blockSize )} ...' )
168136
169137 for idx , path in enumerate (chunk_image_paths ):
170- if idx > 0 :
138+ if idx < len ( chunk_image_paths ) - 1 :
171139 with torch .no_grad ():
172- image1 , h_ori , w_ori , pixelAspectRatio , orientation = image .loadImage (str (chunk_image_paths [idx - 1 ]), True )
173- image2 , h_ori , w_ori , pixelAspectRatio , orientation = image .loadImage (str (chunk_image_paths [idx ]), True )
140+ image1 , h_ori , w_ori , pixelAspectRatio , orientation = image .loadImage (str (chunk_image_paths [idx ]), True )
141+ image2 , h_ori , w_ori , pixelAspectRatio , orientation = image .loadImage (str (chunk_image_paths [idx + 1 ]), True )
174142
143+ scale = 1.0
175144 if max (h_ori , w_ori ) > chunk .node .maxImageSize .value :
176145 maxDim = float (max (h_ori , w_ori ))
177146 scale = float (chunk .node .maxImageSize .value ) / maxDim
@@ -197,34 +166,36 @@ def processChunk(self, chunk):
197166 flow_low , flow_up = model (image1 , image2 , iters = chunk .node .IterationNumber .value , test_mode = True )
198167 flow_up = padder .unpad (flow_up )
199168 flow_up = flow_up [0 ].permute (1 ,2 ,0 ).cpu ().numpy ()
200- flow_out = flow_up .copy ()
201-
202- h ,w ,_ = flow_out .shape
203- nc = np .zeros ((h ,w ,1 ), dtype = flow_out .dtype )
204- flow_out = np .concatenate ((flow_out ,nc ), axis = 2 )
169+ flow_up_out = flow_up .copy () / scale
170+ flow_low = padder .unpad (flow_low )
171+ flow_low = flow_low [0 ].permute (1 ,2 ,0 ).cpu ().numpy ()
172+ flow_low_out = flow_low .copy () / scale
173+
174+ h ,w ,_ = flow_up_out .shape
175+ nc = np .zeros ((h ,w ,1 ), dtype = flow_up_out .dtype )
176+ flow_up_out = np .concatenate ((flow_up_out ,nc ), axis = 2 )
177+ h ,w ,_ = flow_low_out .shape
178+ nc = np .zeros ((h ,w ,1 ), dtype = flow_low_out .dtype )
179+ flow_low_out = np .concatenate ((flow_low_out ,nc ), axis = 2 )
205180
206181 outputDirPath = Path (chunk .node .output .value )
207182 image_stem = Path (chunk_image_paths [idx ]).stem
208- of_file_name = image_stem + ".exr"
183+ ofup_file_name = image_stem + ".exr"
209184
210- image .writeImage (str (outputDirPath / of_file_name ), flow_out , h_ori , w_ori , orientation , pixelAspectRatio )
185+ image .writeImage (str (outputDirPath / ofup_file_name ), flow_up_out , h_ori , w_ori , orientation , pixelAspectRatio )
211186
212- chunk .logger .info ('Publish end' )
187+ chunk .logger .info ('RaftOF end' )
213188 finally :
214189 chunk .logManager .end ()
215190
216- def get_image_paths_list (input_path , extension ):
191+ def get_image_paths_list (input_path ):
217192 from pyalicevision import sfmData
218193 from pyalicevision import sfmDataIO
219194 from pathlib import Path
220- import itertools
221195
222- include_suffixes = [extension .lower (), extension .upper ()]
223196 image_paths = []
224197
225- if Path (input_path ).is_dir ():
226- image_paths = sorted (itertools .chain (* (Path (input_path ).glob (f'*.{ suffix } ' ) for suffix in include_suffixes )))
227- elif Path (input_path ).suffix .lower () in [".sfm" , ".abc" ]:
198+ if Path (input_path ).suffix .lower () in [".sfm" , ".abc" ]:
228199 if Path (input_path ).exists ():
229200 dataAV = sfmData .SfMData ()
230201 if sfmDataIO .load (dataAV , input_path , sfmDataIO .ALL ):
0 commit comments