Big Data: Parallel Processing with Apache Beam and GCP Dataflow

def pl__3__parallel_download_videos(vid_index_schemad_pcoll, d_pl_options, n_partitions=8):
  vid_index_schemad_pcoll_download_partitions = (
    vid_index_schemad_pcoll
    | "Beam PL: partition schemad video index for download parallelization" >> beam.Partition(
        lambda vid_index_row, num_partitions: random.randint(0,num_partitions-1), 
        n_partitions
      )
  )
  partition_download_results = [None for i in range(n_partitions)]
  for i, vid_index_schemad_pcoll_partition in enumerate(vid_index_schemad_pcoll_download_partitions):
    p_label = f"p{i+1}"
    p_label_indented = f"\t{p_label}"
    p_dl_results = (
      vid_index_schemad_pcoll_partition
      | f"Beam PL: {p_label} gather download info for video segments" >> beam.ParDo(VideoSegmentInfoGatherer(d_pl_options))
      | f"Beam PL: {p_label} download video segments" >> beam.ParDo(VideoSegmentDownloader(d_pl_options, f"{p_label_indented}"))
    )
    partition_download_results[i] = p_dl_results
  merged_download_results = (
    (p_dl_r for p_dl_r in partition_download_results) 
    | f"Beam PL: merge download results" >> beam.Flatten() 
  )
  return merged_download_results

def pl__4__parallel_extract_target_video_frames(merged_download_results, d_pl_options, n_partitions=8):
  """
  # ******************** EXTRACT SEGMENT-FRAMES IN PARALLEL: BEGIN ********************
  #   NOTE! THIS IS A CRUCIAL PIECE SO PAY ATTENTION TO THE FOLLOWING!!
  #   ********** --> IMPORTANT VIDEO-FRAME EXTRACTION PROCESSING INFORMATION<-- (BEGIN) **********
  #     We partitioned vid_index_schemad_pcoll so that video-SEGMENT downloads can occur independently.
  #     Downloading segments can occur independently since there is no correlation between each segment
  #       AS FAR AS DOWNLOADING IS CONCERNED.
  #
  #     However, AS FAR AS EXTRACTION IS CONCERNED, each segment is related by the target video composed
  #       of each segment.  The segment-videos themselves are ordered as they compose the final target
  #       video corresponding of ordered segment videos. For example, if a target video is composed of
  #       three segment videos, those segments occur in a certain order, as specified by the video index.
  #       Expanding upon this example, suppose target video "some_story_given_by_john_doe_0.mov", was recorded
  #       and saved in three corresponding video segments (to save space, I guess?) 
  #       "some_story_given_by_john_doe_0_1.mov", "some_story_given_by_john_doe_0_2.mov", and
  #       "some_story_given_by_john_doe_0_3.mov". Note that the trailing "0" in the TARGET VIDEO filename
  #       indicates the camera perspective... all stories are potentially filmed from multiple synchronized
  #       camera perspectives/angles - there were obvioiusly multiple synchronized video recorders used in
  #       in that case.  However, for this example, we are focusing on the target video for camera perspective 0.
  #       Anyway, as said, there are three segments which compose the target video.  THESE SEGMENT VIDEOS
  #       ARE ORDERED (in time).  THEREFORE, THE FRAMES COMPOSING EACH SEGMENT VIDEO ARE CONSEQUENTLY ORDERED
  #       (in time).  THE BOTTOM LINE IS THAT WE NOW NEED TO GROUP SEGMENT VIDEOS, KEYED BY CORRESPONDING
  #       TARGET VIDEO.  FURTHERMORE, THE COLLECTION OF SEGMENT VIDEOS FOR EACH TARGET VIDEO MUST BE ORDERED.
  #       THAT IS, WE MUST EXTRACT SEGMENT FRAMES AND SAVE THEM TO THE FILE SYSTEM WITH A FILE NAMING SCHEME
  #       THAT REFLECTS FRAME ORDER OF THE UNION OF ALL SEGMENT FRAMES.  IF WE EXTRACT THE FRAMES OF EACH
  #       ORDERED SEGMENT, THEN A SIMPLE NUMERIC INDEX AS SEGMENT-FRAME FILENAME WILL DO THE TRICK.
  #   ********** --> IMPORTANT VIDEO-FRAME EXTRACTION PROCESSING INFORMATION<-- (END) **********
  """
  # GROUP segment videos by target video
  #   note that this depends on the DAG - i.e. will not occur until partition_download_results are ready which, of course, does not occur until all videos have been downloaded
  target_vid_seg_frame_extraction_partitions = (
    merged_download_results
    | f"Beam PL: group extraction info for video segments by target video" >> beam.GroupBy(lambda d: d['target_video_fname'])
    | f"Beam PL: partition target video segment info for extraction parallelization" >> beam.Partition(
        lambda vid_index_row, num_partitions: random.randint(0,num_partitions-1), 
        n_partitions
      )
  )
  partition_extraction_results = [None for i in range(n_partitions)]
  for i, p in enumerate(target_vid_seg_frame_extraction_partitions):
    p_label = f"p{i+1}"
    p_label_indented = f"\t{p_label}"
    p_extraction_results = (
      p
      | f"Beam PL: {p_label} extract frames of each segment per target video" >> beam.ParDo(SegmentFrameExtractor(d_pl_options, f"{p_label_indented}", debug=False))
    )
    partition_extraction_results[i] = p_extraction_results
    (
      p_extraction_results
      | f"Beam PL: {p_label} count target videos processed" >> beam.combiners.Count.Globally() 
      | f"Beam PL: {p_label} print target videos processed count" >> beam.ParDo(beam__common.PipelinePcollPrinter(label=p_label_indented, msg="target videos processed"))
    )
  merged_extraction_results = (
    (p_extraction_results for p_extraction_results in partition_extraction_results) 
    | f"Beam PL: merge extraction results" >> beam.Flatten()
  )
  _ = (
    merged_extraction_results
    | "Beam PL: apply schema to merged extraction results pcoll" >> beam.Map(lambda x: beam.Row(
          video_fname=str(x[0]),
          n_stitched_frames=int(x[1])
        ))
    # | "Beam PL: count total frames extracted" >> beam.transforms.sql.SqlTransform(f"SELECT SUM(n_stitched_frames) AS total_frames_extracted FROM PCOLLECTION") # this is VERY, VERY SLOW
    | "Beam PL: select n_stitched_frames" >> beam.Map(lambda extraction_results_row: extraction_results_row.n_stitched_frames)
    | "Beam PL: count total frames extracted" >> beam.CombineGlobally(sum)
    | f"Beam PL: print total frames extracted" >> beam.ParDo(beam__common.PipelinePcollPrinter(msg="TOTAL FRAMES EXTRACTED"))
  )
  return merged_extraction_results

def beam_extract_frames(tpl_target_video_extraction_info, d_pl_options, label="", debug=False):
  """
  expects tpl_target_video_extraction_info: (video_fname, list({'target_video_fname': target_video_fname, 'target_video_frames_dir': target_video_frames_dir, 'segment_url': str(url), 'segment_fname': str(url).split('/')[-1]}))
  """

  # # log_results = []
  target_video_fname = tpl_target_video_extraction_info[0]
  segment_dicts = sorted(tpl_target_video_extraction_info[1], key=lambda segment_dict: segment_dict['segment_fname'])
  target_video_frames_dir = segment_dicts[0]['target_video_frames_dir']

  target_stitched_vid_name = target_video_frames_dir.split(os.path.sep)[-1]
  if not fileio.dir_path_exists(target_video_frames_dir, d_pl_options)[0]:
    fileio.make_dirs(target_video_frames_dir, d_pl_options)

  video_dir = d_pl_options[fidscs_globals.OPT_NAME_VIDEO_DIR]
  local_vid_segment_paths = [fileio.path_join(video_dir, segment_dict['segment_fname']) for segment_dict in segment_dicts]
  for segment_dict in segment_dicts:
    segment_dict['n_frames_extracted'] = 0

  # create local dir for extraction (since OpenCV works only with local file system currently) if we have GCS filesystem
  truly_local_vid_dir = None
  truly_local_vid_dir_suffix = None
  fs = FileSystems.get_filesystem(video_dir)
  if type(fs) == GCSFileSystem:
    truly_local_vid_dir_suffix = '/'.join(video_dir.split('/')[1:])
    truly_local_vid_dir = '/tmp'+truly_local_vid_dir_suffix
    # print(f"\t\tGCS storage detected! Extracting frames to truly_local_vid_dir {truly_local_vid_dir} (and will then upload to GCS after that)...")
    if debug: print(f"\t\t{truly_local_vid_dir} exists: {fileio.dir_path_exists(truly_local_vid_dir, d_pl_options)}")
    if not fileio.dir_path_exists(truly_local_vid_dir, d_pl_options)[0]:
      if debug: print(f"\tcreating {truly_local_vid_dir}...")
      truly_local_vid_dir_path_segs = truly_local_vid_dir.split('/')
      if debug: print(f"\t\ttruly_local_vid_dir_path_segs: {truly_local_vid_dir_path_segs}")
      s_cum_path = ''
      for i, truly_local_vid_dir_path_seg in enumerate(truly_local_vid_dir_path_segs[1:]):
        s_cum_path += '/'+truly_local_vid_dir_path_seg
        fileio.make_dirs(s_cum_path, d_pl_options)
      if debug: print(f"\t\t{s_cum_path} exists: {fileio.dir_path_exists(s_cum_path, d_pl_options)}")

  vc_results = [capture_segment_video(local_vid_segment_path, truly_local_vid_dir, d_pl_options, debug=debug) for local_vid_segment_path in local_vid_segment_paths]
  vid_caps = [vc_result[0] for vc_result in vc_results]
  truly_local_target_video_frames_dirs = [vc_result[1] for vc_result in vc_results]

  for seg_vid_cap in vid_caps:
    seg_vid_cap.set(cv2.CAP_PROP_FPS, fidscs_globals.FPS)
  frame_counts = list(map(lambda vc: int(vc.get(cv2.CAP_PROP_FRAME_COUNT)), vid_caps))
  n_frames_expected = sum(frame_counts)

  failed_target_videos = []

  n_stitched_frames = 0
  if n_frames_expected > 0:
    # get count of existing stitched frames in target_stitched_vid_frames_dir
    n_stitched_frames = len(fileio.list_dir(target_video_frames_dir, d_pl_options))

    b_restitch = n_stitched_frames < n_frames_expected
    n_stitched_frames = 0 if b_restitch else n_stitched_frames

    for i, seg_vid_cap in enumerate(vid_caps):
      segment_dict = segment_dicts[i]
      _n_frames_expected = frame_counts[i]

      if b_restitch:
        success, frame = seg_vid_cap.read()
        n_frames = 0
        while success:
          write_frame_to_file(
            frame, 
            n_stitched_frames, 
            target_video_frames_dir, 
            truly_local_target_video_frames_dir=truly_local_target_video_frames_dirs[i], 
            debug=debug
          )

          n_frames += 1
          n_stitched_frames += 1
          success, frame = seg_vid_cap.read()

        seg_path = local_vid_segment_paths[i]
        seg_fname = seg_path.split(os.path.sep)[-1]
        if n_frames != _n_frames_expected:
          print(f"{label+': ' if len(label)>0 else ''}{fidscs_globals.VALIDATION_FATAL_ERROR_TEXT} Cannot stitch together target video {target_video_fname} since {_n_frames_expected} frames were expected from segment {seg_fname} ({seg_path}) but only {n_frames} were successfully extracted")
          failed_target_videos.append(target_video_fname)
          fail = True
          break
        else:
          print(f"{label+': ' if len(label)>0 else ''}Added {n_stitched_frames} frames from segment {seg_fname} for target video {target_video_fname} (stitched-frames dir {target_video_frames_dir})")

      else:
        n_frames = _n_frames_expected
        print(f"{label+': ' if len(label)>0 else ''}Found existing stiched-frames for {target_stitched_vid_name} ({n_stitched_frames} frames in {target_video_frames_dir})")

      segment_dict['n_frames_extracted'] = n_frames

  else:
    if fidscs_globals.OUTPUT_INFO_LEVEL <= fidscs_globals.OUTPUT_INFO_LEVEL__WARNING:
      print(f"\t{fidscs_globals.VALIDATION_WARNING_TEXT} Cannot stitch together target video {target_video_fname} since cv2.CAP_PROP_FRAME_COUNT reports segments have zero frames")
    failed_target_videos.append(target_video_fname)
    fail = True

  if truly_local_vid_dir is not None:
    for truly_local_target_video_frames_dir in truly_local_target_video_frames_dirs:
      fileio.delete_file(truly_local_target_video_frames_dir, d_pl_options, recursive=True, debug=True)

  return [(tpl_target_video_extraction_info[0], n_stitched_frames, segment_dicts)]

def capture_segment_video(vid_segment_path, truly_local_vid_dir, d_pl_options, debug=False):
  video_fname = vid_segment_path.split('/')[-1]

  truly_local_target_video_frames_dir = None

  fs = FileSystems.get_filesystem(vid_segment_path)
  if type(fs) == GCSFileSystem:
    if debug: print(f"\n\n\tattempting to open video {vid_segment_path} for reading...")
    with fileio.open_file_read(vid_segment_path) as f:
      if debug: print(f"\t\tSUCCESS")

      # now read from local bytes and write to GCS
      buffer = f.read()
      truly_local_vid_segment_path = truly_local_vid_dir+'/'+video_fname
      if debug: print(f"\t\tattempting to write {truly_local_vid_segment_path} (truly) locally...")
      with fileio.open_file_write(truly_local_vid_segment_path) as f_local:
        f_local.write(buffer)
        f_local.close()
        if debug: print(f"\t\t\tSUCCESS")
      f.close()

      vid_segment_path = truly_local_vid_segment_path

      # (truly local) dir for saving frames
      truly_local_target_video_frames_dir = truly_local_vid_dir+'/'+fidscs_globals.STICHED_VIDEO_FRAMES_DIR_NAME+'/'+video_fname.split('.')[0]
      if debug: print(f"\t\t\tattempting to create directory {truly_local_target_video_frames_dir} (truly_local_target_video_frames_dir) for frames extracted from (truly local) video {truly_local_vid_segment_path}...")
      if not fileio.dir_path_exists(truly_local_target_video_frames_dir, d_pl_options)[0]:
        if debug: print(f"\t\t\t\tcreating {truly_local_target_video_frames_dir}...")
        fileio.make_dirs(truly_local_target_video_frames_dir, d_pl_options)
      truly_local_target_video_frames_dir_exists = fileio.dir_path_exists(truly_local_target_video_frames_dir, d_pl_options)[0]
      if debug: print(f"\t\t\t\t\t{truly_local_target_video_frames_dir} exists: {truly_local_target_video_frames_dir_exists}")
      if not truly_local_target_video_frames_dir_exists:
        raise Exception(f"required directory truly_local_target_video_frames_dir {truly_local_target_video_frames_dir_exists} does not exist")

  if debug: print(f"\t\t\tattempting to capture (cv2.VideoCapture) video {vid_segment_path})...")

  # finally, capture the video bytes
  return cv2.VideoCapture(vid_segment_path), truly_local_target_video_frames_dir

def write_frame_to_file(frame, index, target_video_frames_dir, truly_local_target_video_frames_dir=None, debug=False):
  local_frame_path = fileio.path_join(target_video_frames_dir, f"{index}.jpg") # this is the final frame path

  if truly_local_target_video_frames_dir is not None:
    # write truly local frame file
    truly_local_frame_path = truly_local_target_video_frames_dir+'/'+f"{index}.jpg"
    if debug: print(f"\t\t\t\t\t\tattempting to write {truly_local_frame_path} frame...")
    cv2.imwrite(truly_local_frame_path, frame)
    if debug: print(f"\t\t\t\t\t\t\tSUCCESS")
    if debug: print(f"\t\t\t\t\t\t\tattempting to open {truly_local_frame_path} for read...")
    with fileio.open_file_read(truly_local_frame_path) as f_truly_local_frame:
      buffer = f_truly_local_frame.read()
      if debug: print(f"\t\t\t\t\t\t\t\tSUCCESS")
      if debug: print(f"\t\t\t\t\t\t\t\t\tattempting to open {local_frame_path} for final write...")
      with fileio.open_file_write(local_frame_path) as f_frame_final:
        f_frame_final.write(buffer)
        f_frame_final.close()
        if debug: print(f"\t\t\t\t\t\t\t\t\t\tSUCCESS")
      buffer = None
      f_truly_local_frame.close()

  else:
    if debug: print(f"\t\t\t\t\t\t\t\t\tattempting to open {local_frame_path} for final write...")
    cv2.imwrite(local_frame_path, frame)
    if debug: print(f"\t\t\t\t\t\t\t\t\t\tSUCCESS")