vace/vace_preproccess.py

# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

import os
import copy
import time
import inspect
import argparse
import importlib

from configs import VACE_PREPROCCESS_CONFIGS
import annotators
from annotators.utils import read_image, read_mask, read_video_frames, save_one_video, save_one_image


def parse_bboxes(s):
    bboxes = []
    for bbox_str in s.split():
        coords = list(map(float, bbox_str.split(',')))
        if len(coords) != 4:
            raise ValueError(f"The bounding box requires 4 values, but the input is {len(coords)}.")
        bboxes.append(coords)
    return bboxes

def validate_args(args):
    assert args.task in VACE_PREPROCCESS_CONFIGS, f"Unsupport task: [{args.task}]"
    assert args.video is not None or args.image is not None or args.bbox is not None, "Please specify the video or image or bbox."
    return args

def get_parser():
    parser = argparse.ArgumentParser(
        description="Data processing carried out by VACE"
    )
    parser.add_argument(
        "--task",
        type=str,
        default='',
        choices=list(VACE_PREPROCCESS_CONFIGS.keys()),
        help="The task to run.")
    parser.add_argument(
        "--video",
        type=str,
        default=None,
        help="The path of the videos to be processed, separated by commas if there are multiple.")
    parser.add_argument(
        "--image",
        type=str,
        default=None,
        help="The path of the images to be processed, separated by commas if there are multiple.")
    parser.add_argument(
        "--mode",
        type=str,
        default=None,
        help="The specific mode of the task, such as firstframe, mask, bboxtrack, label...")
    parser.add_argument(
        "--mask",
        type=str,
        default=None,
        help="The path of the mask images to be processed, separated by commas if there are multiple.")
    parser.add_argument(
        "--bbox",
        type=parse_bboxes,
        default=None,
        help="Enter the bounding box, with each four numbers separated by commas (x1, y1, x2, y2), and each pair separated by a space."
    )
    parser.add_argument(
        "--label",
        type=str,
        default=None,
        help="Enter the label to be processed, separated by commas if there are multiple."
    )
    parser.add_argument(
        "--caption",
        type=str,
        default=None,
        help="Enter the caption to be processed."
    )
    parser.add_argument(
        "--direction",
        type=str,
        default=None,
        help="The direction of outpainting includes any combination of left, right, up, down, with multiple combinations separated by commas.")
    parser.add_argument(
        "--expand_ratio",
        type=float,
        default=None,
        help="The outpainting's outward expansion ratio.")
    parser.add_argument(
        "--expand_num",
        type=int,
        default=None,
        help="The number of frames extended by the extension task.")
    parser.add_argument(
        "--maskaug_mode",
        type=str,
        default=None,
        help="The mode of mask augmentation, such as original, original_expand, hull, hull_expand, bbox, bbox_expand.")
    parser.add_argument(
        "--maskaug_ratio",
        type=float,
        default=None,
        help="The ratio of mask augmentation.")
    parser.add_argument(
        "--pre_save_dir",
        type=str,
        default=None,
        help="The path to save the processed data.")
    parser.add_argument(
        "--save_fps",
        type=int,
        default=16,
        help="The fps to save the processed data.")
    return parser


def preproccess():
    pass

def proccess():
    pass

def postproccess():
    pass

def main(args):
    args = argparse.Namespace(**args) if isinstance(args, dict) else args
    args = validate_args(args)

    task_name = args.task
    video_path = args.video
    image_path = args.image
    mask_path = args.mask
    bbox = args.bbox
    caption = args.caption
    label = args.label
    save_fps = args.save_fps

    # init class
    task_cfg = copy.deepcopy(VACE_PREPROCCESS_CONFIGS)[task_name]
    class_name = task_cfg.pop("NAME")
    input_params = task_cfg.pop("INPUTS")
    output_params = task_cfg.pop("OUTPUTS")

    # input data
    fps = None
    input_data = copy.deepcopy(input_params)
    if 'video' in input_params:
        assert video_path is not None, "Please set video or check configs"
        frames, fps, width, height, num_frames = read_video_frames(video_path.split(",")[0], use_type='cv2',  info=True)
        assert frames is not None, "Video read error"
        input_data['frames'] = frames
        input_data['video'] = video_path
    if 'frames' in input_params:
        assert video_path is not None, "Please set video or check configs"
        frames, fps, width, height, num_frames = read_video_frames(video_path.split(",")[0], use_type='cv2', info=True)
        assert frames is not None, "Video read error"
        input_data['frames'] = frames
    if 'frames_2' in input_params:
        # assert video_path is not None and len(video_path.split(",")[1]) >= 2, "Please set two videos or check configs"
        if  len(video_path.split(",")) >= 2:
            frames, fps, width, height, num_frames = read_video_frames(video_path.split(",")[1], use_type='cv2', info=True)
            assert frames is not None, "Video read error"
            input_data['frames_2'] = frames
    if 'image' in input_params:
        assert image_path is not None, "Please set image or check configs"
        image, width, height = read_image(image_path.split(",")[0], use_type='pil', info=True)
        assert image is not None, "Image read error"
        input_data['image'] = image
    if 'image_2' in input_params:
        # assert image_path is not None and len(image_path.split(",")[1]) >= 2, "Please set two images or check configs"
        if len(image_path.split(",")) >= 2:
            image, width, height = read_image(image_path.split(",")[1], use_type='pil', info=True)
            assert image is not None, "Image read error"
            input_data['image_2'] = image
    if 'images' in input_params:
        assert image_path is not None, "Please set image or check configs"
        images = [ read_image(path, use_type='pil', info=True)[0] for path in image_path.split(",") ]
        input_data['images'] = images
    if 'mask' in input_params:
        # assert mask_path is not None, "Please set mask or check configs"
        if mask_path is not None:
            mask, width, height = read_mask(mask_path.split(",")[0], use_type='pil', info=True)
            assert mask is not None, "Mask read error"
            input_data['mask'] = mask
    if 'bbox' in input_params:
        # assert bbox is not None, "Please set bbox"
        if bbox is not None:
            input_data['bbox'] = bbox[0] if len(bbox) == 1 else bbox
    if 'label' in input_params:
        # assert label is not None, "Please set label or check configs"
        input_data['label'] = label.split(',') if label is not None else None
    if 'caption' in input_params:
        # assert caption is not None, "Please set caption or check configs"
        input_data['caption'] = caption
    if 'mode' in input_params:
        input_data['mode'] = args.mode
    if 'direction' in input_params:
        if args.direction is not None:
            input_data['direction'] = args.direction.split(',')
    if 'expand_ratio' in input_params:
        if args.expand_ratio is not None:
            input_data['expand_ratio'] = args.expand_ratio
    if 'expand_num' in input_params:
        # assert args.expand_num is not None, "Please set expand_num or check configs"
        if args.expand_num is not None:
            input_data['expand_num'] = args.expand_num
    if 'mask_cfg' in input_params:
        # assert args.maskaug_mode is not None and args.maskaug_ratio is not None, "Please set maskaug_mode and maskaug_ratio or check configs"
        if args.maskaug_mode is not None:
            if args.maskaug_ratio is not None:
                input_data['mask_cfg'] = {"mode": args.maskaug_mode, "kwargs": {'expand_ratio': args.maskaug_ratio, 'expand_iters': 5}}
            else:
                input_data['mask_cfg'] = {"mode": args.maskaug_mode}

    # processing
    pre_ins = getattr(annotators, class_name)(cfg=task_cfg, device=f'cuda:{os.getenv("RANK", 0)}')
    results = pre_ins.forward(**input_data)

    # output data
    save_fps = fps if fps is not None else save_fps
    if args.pre_save_dir is None:
        pre_save_dir = os.path.join('processed', task_name, time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())))
    else:
        pre_save_dir = args.pre_save_dir
    if not os.path.exists(pre_save_dir):
        os.makedirs(pre_save_dir)

    ret_data = {}
    if 'frames' in output_params:
        frames =  results['frames'] if isinstance(results, dict) else results
        if frames is not None:
            save_path = os.path.join(pre_save_dir, f'src_video-{task_name}.mp4')
            save_one_video(save_path, frames, fps=save_fps)
            print(f"Save frames result to {save_path}")
            ret_data['src_video'] = save_path
    if 'masks' in output_params:
        frames = results['masks'] if isinstance(results, dict) else results
        if frames is not None:
            save_path = os.path.join(pre_save_dir, f'src_mask-{task_name}.mp4')
            save_one_video(save_path, frames, fps=save_fps)
            print(f"Save frames result to {save_path}")
            ret_data['src_mask'] = save_path
    if 'image' in output_params:
        ret_image =  results['image'] if isinstance(results, dict) else results
        if ret_image is not None:
            save_path = os.path.join(pre_save_dir, f'src_ref_image-{task_name}.png')
            save_one_image(save_path, ret_image, use_type='pil')
            print(f"Save image result to {save_path}")
            ret_data['src_ref_images'] = save_path
    if 'images' in output_params:
        ret_images = results['images'] if isinstance(results, dict) else results
        if ret_images is not None:
            src_ref_images = []
            for i, img in enumerate(ret_images):
                if img is not None:
                    save_path = os.path.join(pre_save_dir, f'src_ref_image_{i}-{task_name}.png')
                    save_one_image(save_path, img, use_type='pil')
                    print(f"Save image result to {save_path}")
                    src_ref_images.append(save_path)
            if len(src_ref_images) > 0:
                ret_data['src_ref_images'] = ','.join(src_ref_images)
            else:
                ret_data['src_ref_images'] = None
    if 'mask' in output_params:
        ret_image =  results['mask'] if isinstance(results, dict) else results
        if ret_image is not None:
            save_path = os.path.join(pre_save_dir, f'src_mask-{task_name}.png')
            save_one_image(save_path, ret_image, use_type='pil')
            print(f"Save mask result to {save_path}")
    return ret_data


if __name__ == "__main__":
    args = get_parser().parse_args()
    main(args)