scripts/quick_utils.py

import json
import cv2
import numpy as np
from pathlib import Path
import os
from PIL import Image


if __name__ == "__main__":
    '''get frame_id'''
    # with open("/scratch/yuqian_fu/egoexo_val_framelevel_newprompt_all_instruction.json") as f:
    #     data = json.load(f)
    
    # data_new = []
    # for item in data:
    #     if item['video_name'] == "1247a29c-9fda-47ac-8b9c-78b1e76e977e":
    #         data_new.append(item)

    # test_sample = data_new[0]
    # print(test_sample['new_img_id'])

    '''vis_mask'''
    def upsample_mask(mask, frame):
        H, W = frame.shape[:2]
        mH, mW = mask.shape[:2]

        if W > H:
            ratio = mW / W
            h = H * ratio
            diff = int((mH - h) // 2)
            if diff == 0:
                mask = mask
            else:
                mask = mask[diff:-diff]
        else:
            ratio = mH / H
            w = W * ratio
            diff = int((mW - w) // 2)
            if diff == 0:
                mask = mask
            else:
                mask = mask[:, diff:-diff]

        mask = cv2.resize(mask, (W, H))
        return mask

    def blend_mask(input_img, binary_mask, alpha=0.5, color="g"):
        if input_img.ndim == 2:
            return input_img
        mask_image = np.zeros(input_img.shape, np.uint8)
        if color == "r":
            mask_image[:, :, 0] = 255
        if color == "g":
            mask_image[:, :, 1] = 255
        if color == "b":
            mask_image[:, :, 2] = 255
        if color == "o":
            mask_image[:, :, 0] = 255
            mask_image[:, :, 1] = 165
            mask_image[:, :, 2] = 0
        if color == "c":
            mask_image[:, :, 0] = 0
            mask_image[:, :, 1] = 255
            mask_image[:, :, 2] = 255
        if color == "p":
            mask_image[:, :, 0] = 128
            mask_image[:, :, 1] = 0
            mask_image[:, :, 2] = 128

        mask_image = mask_image * np.repeat(binary_mask[:, :, np.newaxis], 3, axis=2)
        blend_image = input_img[:, :, :].copy()
        pos_idx = binary_mask > 0
        for ind in range(input_img.ndim):
            ch_img1 = input_img[:, :, ind]
            ch_img2 = mask_image[:, :, ind]
            ch_img3 = blend_image[:, :, ind]
            ch_img3[pos_idx] = alpha * ch_img1[pos_idx] + (1 - alpha) * ch_img2[pos_idx]
            blend_image[:, :, ind] = ch_img3
        return blend_image

    mask_path = "/scratch/yuqian_fu/test_result/mask/1247a29c-9fda-47ac-8b9c-78b1e76e977e_ref/30_pred_complex_ego_watch.png"
    img_path = "/scratch/yuqian_fu/test_data/1247a29c-9fda-47ac-8b9c-78b1e76e977e/aria01_214-1/30.jpg"
    mask = Image.open(mask_path)
    mask = np.array(mask)
    print(mask.shape)
    mask2 = cv2.imread(mask_path)
    print(type(mask2), mask2.shape)
    frame = cv2.imread(img_path)

    unique_instances = np.unique(mask)
    unique_instances = unique_instances[unique_instances != 0]
    if len(unique_instances) != 0:
        for i,instance in enumerate(unique_instances):
            binary_mask = (mask == instance).astype(np.uint8)
            binary_mask = cv2.resize(binary_mask, (frame.shape[1], frame.shape[0]))
            binary_mask = upsample_mask(binary_mask, frame)
            out = blend_mask(frame, binary_mask, color="g")
            save_path = "/scratch/yuqian_fu/test_result/img/1247a29c-9fda-47ac-8b9c-78b1e76e977e_ref/30_pred_complex_ego_watch.jpg"
            Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)
            cv2.imwrite(save_path, out)


    '''change insttruction'''
    # with open("/scratch/yuqian_fu/egoexo_val_framelevel_newprompt_all_instruction.json") as f:
    #     data = json.load(f)
    
    # data_new = []
    # for item in data:
    #     if item['video_name'] == "1247a29c-9fda-47ac-8b9c-78b1e76e977e":
    #         data_new.append(item)

    # test_sample = data_new[0]
    # # print(test_sample['new_img_id'])
    # # print(test_sample['image'])
    # # print(test_sample['instruction'])

    # instruction_list = []
    # sample = {
    #         "tokens": ['the', 'ball'],
    #         "raw": "the ball.",
    #         "sent_id": 2203,
    #         "sent": "the ball"
    # }
    # image_info = {
    #             'file_name': test_sample['first_frame_image'],
    #             'height': 704,
    #             'width': 704,
    #         }
    # instruction_list.append(sample)
    # to_save = {
    #     "image":test_sample['first_frame_image'],
    #     "image_info":image_info,
    #     "anns":test_sample['first_frame_anns'],
    #     "first_frame_image":test_sample['first_frame_image'],
    #     "first_frame_anns":test_sample['first_frame_anns'],
    #     "new_img_id":test_sample['new_img_id'],
    #     "video_name":test_sample['video_name'],
    #     "instruction":instruction_list
    # }
    # save_path = "/scratch/yuqian_fu/sample_instruction_ego.json"
    # with open(save_path, "w") as f:
    #     json.dump([to_save], f)