DataFlow appears to be stuck. - on reading images from Google Storage - tensorflow

Question

I have a DataFlow job, which first reads in 2 text files, located in Google Cloud Storage. The text files contain the paths to images also located in Google Cloud Storage.
After some inspections, I can confirm that reading the text files is successful, but the DataFlow job is stuck at reading the images. The same code runs perfectly locally. Which makes me think that maybe the image paths are incorrect, but it's not.
Here's my job ID: 2018-01-10_12_16_56-8294573519126715750
Any advice would be appreciated. Also any pointers on how to solve / debug this problem would be really useful as I don't even know where to start.
Thanks
Pipeline Definition
def configure_pipeline(pipeline, args):
read_input_source = beam.io.ReadFromText(args.input_path, strip_trailing_newlines=True)
read_img_paths = beam.io.ReadFromText(args.input_imgs, strip_trailing_newlines=True)
img_paths = (pipeline | 'Read image paths' >> read_img_paths)
train_points = (pipeline | 'Read data point' >> read_input_source)
_ = (train_points | "Read image" >> beam.ParDo(ExtractDataDoFn(), beam.pvalue.AsIter(img_paths)))
Read Images - Most of the code is simply parsing the image paths from the text file and some data to index the image.
class ExtractDataDoFn(beam.DoFn):
def start_bundle(self, context=None):
# Each frame has its own path to its image
self.frame_number_to_name = {}
def process(self, element, img_paths):
try:
line = element.element
except AttributeError:
pass
if not self.frame_number_to_name:
for path in img_paths:
if len(path) > 4:
frame_number = int(path[-10 : -4])
self.frame_number_to_name[frame_number] = path
line_tokens = element.split(':')
pivot_example = line_tokens[0].strip('\'')
example = line_tokens[1].strip('\'')
label = int(line_tokens[2])
# Get image paths
pivot_frame_number = int(pivot_example.split(',')[0])
pivot_path = self.frame_number_to_name[pivot_frame_number]
example_frame_number = int(example.split(',')[0])
example_path = self.frame_number_to_name[example_frame_number]
# Read images
def _open_file_read_binary(uri):
try:
return file_io.FileIO(uri, mode='rb')
except errors.InvalidArgumentError:
return file_io.FileIO(uri, mode='r')
# Read pivot
try:
with _open_file_read_binary(pivot_path) as f:
pivot_image_bytes = f.read()
pivot_img = Image.open(io.BytesIO(pivot_image_bytes)).convert('RGB')
except Exception as e: # pylint: disable=broad-except
logging.exception('Error processing image %s: %s', pivot_example, str(e))
return
# Read example
try:
with _open_file_read_binary(example_path) as f:
example_image_bytes = f.read()
example_img = Image.open(io.BytesIO(example_image_bytes)).convert('RGB')
except Exception as e: # pylint: disable=broad-except
logging.exception('Error processing image %s: %s', example, str(e))
return
# Convert to Numpy array
pivot_np = np.array(pivot_img)
example_np = np.array(example_img)
def _get_feature(line, img):
frame_number = int(line.split(',')[0])
y, x = int(line.split(',')[3]), int(line.split(',')[2])
h, w = int(line.split(',')[5]), int(line.split(',')[4])
bb = img[y : y + h, x : x + w, :]
return bb
# Get raw content of bounding box
pivot_feature = _get_feature(pivot_example, pivot_np)
example_feature = _get_feature(example, example_np)
# Resize data
pivot_feature = Image.fromarray(pivot_feature).resize((224, 224))
example_feature = Image.fromarray(example_feature).resize((224, 224))
# Convert back to numpy
pivot_feature = np.array(pivot_feature, np.float64)
example_feature = np.array(example_feature, np.float64)
# print(pivot_feature.shape)
yield pivot_feature, example_feature, label