Skip to content

Commit b68356d

Browse files
committed
Dynamic size updates
1 parent e1e0eaf commit b68356d

2 files changed

Lines changed: 36 additions & 103 deletions

File tree

meshroom/DenseMotion/DenseMotion.py

Lines changed: 7 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,8 @@
1-
__version__ = "1.0"
1+
__version__ = "2.0"
22

33
from meshroom.core import desc
44
from meshroom.core.utils import VERBOSE_LEVEL
5-
6-
7-
class DenseMotionNodeSize(desc.MultiDynamicNodeSize):
8-
def computeSize(self, node):
9-
from pathlib import Path
10-
import itertools
11-
12-
input_path_param = node.attribute(self._params[0])
13-
extension_param = node.attribute(self._params[1])
14-
15-
input_path = input_path_param.value
16-
extension = extension_param.value
17-
include_suffixes = [extension.lower(), extension.upper()]
18-
19-
size = 1
20-
if Path(input_path).is_dir():
21-
image_paths = list(itertools.chain(*(Path(input_path).glob(f'*.{suffix}') for suffix in include_suffixes)))
22-
size = len(image_paths)
23-
elif node.attribute(self._params[0]).isLink:
24-
size = node.attribute(self._params[0]).getLinkParam().node.size
25-
26-
return size
27-
5+
from pyalicevision import parallelization as avpar
286

297
class DenseMotionBlockSize(desc.Parallelization):
308
def getSizes(self, node):
@@ -45,25 +23,16 @@ class DenseMotion(desc.Node):
4523

4624
gpu = desc.Level.INTENSIVE
4725

48-
size = DenseMotionNodeSize(['inputImages', 'inputExtension'])
26+
size = avpar.DynamicViewsSize("inputImages")
4927
parallelization = DenseMotionBlockSize()
5028

5129
inputs = [
5230
desc.File(
5331
name="inputImages",
5432
label="Input Images",
55-
description="Input images to estimate the depth from. Folder path or sfmData filepath",
33+
description="Input images to estimate the depth from, sfmData filepath",
5634
value="",
5735
),
58-
desc.ChoiceParam(
59-
name="inputExtension",
60-
label="Input Extension",
61-
description="Extension of the input images. This will be used to determine which images are to be used if \n"
62-
"a directory is provided as the input. If \"\" is selected, the provided input will be used as such.",
63-
values=["jpg", "jpeg", "png", "exr"],
64-
value="exr",
65-
exclusive=True,
66-
),
6736
desc.FloatParam(
6837
name="PyramidScale",
6938
label="Pyramid Scale",
@@ -98,7 +67,6 @@ class DenseMotion(desc.Node):
9867
value=50,
9968
description="Sets the number of images to process in one chunk. If set to 0, all images are processed at once.",
10069
range=(0, 1000, 1),
101-
enabled=lambda node: node.model.value == "MoGe"
10270
),
10371
desc.ChoiceParam(
10472
name="verboseLevel",
@@ -122,15 +90,13 @@ class DenseMotion(desc.Node):
12290
description="Output optical flow images",
12391
semantic="image",
12492
value=lambda attr: "{nodeCacheFolder}/<FILESTEM>.exr",
125-
group="",
12693
)
12794
]
12895

12996
def preprocess(self, node):
130-
extension = node.inputExtension.value
13197
input_path = node.inputImages.value
13298

133-
image_paths = get_image_paths_list(input_path, extension)
99+
image_paths = get_image_paths_list(input_path)
134100

135101
if len(image_paths) == 0:
136102
raise FileNotFoundError(f'No image files found in {input_path}')
@@ -184,18 +150,14 @@ def processChunk(self, chunk):
184150
finally:
185151
chunk.logManager.end()
186152

187-
def get_image_paths_list(input_path, extension):
153+
def get_image_paths_list(input_path):
188154
from pyalicevision import sfmData
189155
from pyalicevision import sfmDataIO
190156
from pathlib import Path
191-
import itertools
192157

193-
include_suffixes = [extension.lower(), extension.upper()]
194158
image_paths = []
195159

196-
if Path(input_path).is_dir():
197-
image_paths = sorted(itertools.chain(*(Path(input_path).glob(f'*.{suffix}') for suffix in include_suffixes)))
198-
elif Path(input_path).suffix.lower() in [".sfm", ".abc"]:
160+
if Path(input_path).suffix.lower() in [".sfm", ".abc"]:
199161
if Path(input_path).exists():
200162
dataAV = sfmData.SfMData()
201163
if sfmDataIO.load(dataAV, input_path, sfmDataIO.ALL):

meshroom/DenseMotion/RaftOF.py

Lines changed: 29 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,9 @@
1-
__version__ = "1.0"
1+
__version__ = "2.0"
22

33
from re import M
44
from meshroom.core import desc
55
from 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

309
class 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

Comments
 (0)