48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
from pocketflow import Flow, BatchFlow
|
|
from nodes import LoadImage, ApplyFilter, SaveImage
|
|
|
|
def create_base_flow():
|
|
"""Create the base Flow for processing a single image."""
|
|
# Create nodes
|
|
load = LoadImage()
|
|
filter_node = ApplyFilter()
|
|
save = SaveImage()
|
|
|
|
# Connect nodes
|
|
load - "apply_filter" >> filter_node
|
|
filter_node - "save" >> save
|
|
|
|
# Create and return flow
|
|
return Flow(start=load)
|
|
|
|
class ImageBatchFlow(BatchFlow):
|
|
"""BatchFlow for processing multiple images with different filters."""
|
|
|
|
def prep(self, shared):
|
|
"""Generate parameters for each image-filter combination."""
|
|
# List of images to process
|
|
images = ["cat.jpg", "dog.jpg", "bird.jpg"]
|
|
|
|
# List of filters to apply
|
|
filters = ["grayscale", "blur", "sepia"]
|
|
|
|
# Generate all combinations
|
|
params = []
|
|
for img in images:
|
|
for f in filters:
|
|
params.append({
|
|
"input": img,
|
|
"filter": f
|
|
})
|
|
|
|
return params
|
|
|
|
def create_flow():
|
|
"""Create the complete batch processing flow."""
|
|
# Create base flow for single image processing
|
|
base_flow = create_base_flow()
|
|
|
|
# Wrap in BatchFlow for multiple images
|
|
batch_flow = ImageBatchFlow(start=base_flow)
|
|
|
|
return batch_flow |