mbench/gpt_ref-ytvos-revised.py

import sys
from os import path as osp
sys.path.append(osp.abspath(osp.join(osp.dirname(__file__), '..')))

from mbench.ytvos_ref import build as build_ytvos_ref
import argparse
import opts

import sys
from pathlib import Path
import os
from os import path as osp
import skimage
from io import BytesIO

import numpy as np
import pandas as pd
import regex as re
import json

import cv2
from PIL import Image, ImageDraw
import torch
from torchvision.transforms import functional as F

from skimage import measure                        # (pip install scikit-image)
from shapely.geometry import Polygon, MultiPolygon # (pip install Shapely)

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle


import ipywidgets as widgets
from IPython.display import display, clear_output

from openai import OpenAI
import base64

# Function to encode the image
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# Captioner
ytvos_category_valid_list = [
    'airplane', 'ape', 'bear', 'bike', 'bird', 'boat', 'bus', 'camel', 'cat', 'cow', 'crocodile', 
    'deer', 'dog', 'dolphin', 'duck', 'eagle', 'earless_seal', 'elephant', 'fish', 'fox', 'frog', 
    'giant_panda', 'giraffe', 'hedgehog', 'horse', 'leopard', 'lion', 'lizard', 
    'monkey', 'motorbike', 'mouse', 'owl', 'parrot', 'penguin', 'person', 
    'rabbit', 'raccoon', 'sedan', 'shark', 'sheep', 'snail', 'snake',  
    'squirrel', 'tiger', 'train', 'truck', 'turtle', 'whale', 'zebra'
]
def getCaption(video_id, json_data):
    #데이터 가져오기
    video_data = json_data[video_id]
    frame_names = video_data['frame_names']
    video_path = video_data['video_path']
    
    cat_names = set()
    all_captions = dict()
    for obj_id in list(video_data['annotations'][0].keys()):
        cat_names.add(video_data['annotations'][0][obj_id]['category_name'])

    # cat_names : person, snowboard
    # 1. gpt에서 직접 action의 대상이 될 수 있는가 물어보기
    # 2. ref-youtube-vos 에서 제공하는 카테고리 정보에서 우리가 처리하고 싶은 카테고리 이름만 남긴다

    for cat_name in list(cat_names) :        
        image_paths = [os.path.join(video_path, frame_name + '.jpg') for frame_name in frame_names]
        image_captions = {}

        captioner = OpenAI()

        #0단계: action의 대상이 될 수 있는가?
        is_movable = False  
        if cat_name in ytvos_category_valid_list :
            is_movable = True

        # response_check = captioner.chat.completions.create(
        #     model="gpt-4o",
        #     messages=[
        #         {
        #             "role": "user",
        #             "content": f"""
        #                 Can a {cat_name} be a subject of distinct actions or movements? 
        #                 For example, if {cat_name} is a person, animal, or vehicle, it is likely an action-capable subject. 
        #                 However, if it is an inanimate object like a snowboard, tree, or book, it cannot independently perform actions.
        #                 Respond with YES if {cat_name} can perform distinct actions or movements; otherwise, respond with NONE.
        #                 Answer only YES or NONE.
        #             """
        #         }
        #     ],
        # )
        # response_check_content = response_check.choices[0].message.content.strip().lower()
        # print(f"Movable Check for {cat_name}: {response_check_content}")

        # if response_check_content == "yes": is_movable = True
            
        if not is_movable:
            print(f"Skipping {cat_name}: Determined to be non-movable.")
            continue

        for i in range(len(image_paths)):
            image_path = image_paths[i]
            frame_name = frame_names[i]
            base64_image = encode_image(image_path)
            
            #1단계: 필터링
            #print(f"-----------category name: {cat_name}, frame name: {frame_name}")
            response1 = captioner.chat.completions.create(
                model="chatgpt-4o-latest",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {                            
                                "type": "text",
                                
                                "text": f"""Are there multiple {cat_name}s in the image, each performing distinct and recognizable actions? 
                                        Focus only on clear and prominent actions, avoiding minor or ambiguous ones.
                                        Each action should be unique and clearly associated with a specific object. 
                                        
                                        Respond with YES if:
                                        - The {cat_name}s are people, animals or vehicles, and their actions are distinct and recognizable.
                                        - The {cat_name}s involve clear, distinguishable actions performed independently.

                                        Respond with NONE if:
                                        - The {cat_name}s are objects (e.g., snowboard, tree, books) and do not involve direct interaction with a person.
                                        - Actions are ambiguous, minor, or not clearly visible.

                                        If the {cat_name} is 'snowboard' and it is not actively being used or interacted with by a person, output NONE. 
                                        If the {cat_name} is 'person' and their actions are distinct and clear, output YES.

                                        Answer only YES or NONE."""

                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                            },
                        ],
                    }
                ],
            )
            response_content = response1.choices[0].message.content
            should_caption = True if "yes" in response_content.lower() else False
            #print(f"are {cat_name}s distinguished by action: {response_content}")

            #2단계: dense caption 만들기
            if should_caption:
                response2 = captioner.chat.completions.create(
                    model="chatgpt-4o-latest",
                    messages=[
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                
                                    "text": f"""
                                            Generate a detailed action-centric caption describing the actions of the {cat_name}s in the image. 
                                            1. Focus only on clear, unique, and prominent actions that distinguish each object.
                                            2. Avoid describing actions that are too minor, ambiguous, or not visible from the image.
                                            3. Avoid subjective terms such as 'skilled', 'controlled', or 'focused'. Only describe observable actions.
                                            4. Do not include common-sense or overly general descriptions like 'the elephant walks'.
                                            5. Use dynamic action verbs (e.g., holding, throwing, jumping, inspecting) to describe interactions, poses, or movements.
                                            6. Avoid overly detailed or speculative descriptions such as 'slightly moving its mouth' or 'appears to be anticipating'.
                                            7. Pretend you are observing the scene directly, avoiding phrases like 'it seems' or 'based on the description'.
                                            8. Include interactions with objects or other entities when they are prominent and observable.
                                            9. If the image contains multiple {cat_name}s, describe the actions of each individually and ensure the descriptions are non-overlapping and specific.
                                            Output only the caption.""",
                                },
                                {
                                    "type": "image_url",
                                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                                },
                            ],
                        }
                    ],
                )

                caption = response2.choices[0].message.content
                #print(f"{image_path} - {frame_name}: {caption}")
            else:
                caption = None

            image_captions[frame_name] = caption
        all_captions[cat_name] = image_captions
        
    # final : also prepare valid object ids
    valid_obj_ids = []
    valid_cat_names = list(all_captions.keys())
    for obj_id in list(video_data['annotations'][0].keys()):
        cat = video_data['annotations'][0][obj_id]['category_name']
        if cat in valid_cat_names : valid_obj_ids.append(obj_id)
        
    return all_captions, valid_obj_ids

# Referring expression generator and QA filter
def getRefExp(video_id, frame_name, caption, obj_id, json_data):
    
    # 이미지에 해당 물체 바운딩 박스 그리기
    video_data = json_data[video_id]
    frame_names = video_data['frame_names']
    video_path = video_data['video_path']
    I = skimage.io.imread(osp.join(video_path, frame_name + '.jpg'))
    frame_indx = frame_names.index(frame_name)
    obj_data = video_data['annotations'][frame_indx][obj_id]

    bbox = obj_data['bbox']
    cat_name = obj_data['category_name']
    valid = obj_data['valid']

    if valid == 0:
        print("Object not in this frame!")
        return {}


    x_min, y_min, x_max, y_max = bbox   
    x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)     
    cv2.rectangle(I, (x_min, y_min), (x_max, y_max), (225, 0, 0), 2)
    plt.figure()
    plt.imshow(I)
    plt.axis('off')
    plt.show()
    
    #cropped object for visibility check
    cropped_I = I[y_min:y_max, x_min:x_max]
    pil_cropped_I = Image.fromarray(cropped_I)
    buff_crop = BytesIO()
    pil_cropped_I.save(buff_crop, format='JPEG')
    base64_cropped_I = base64.b64encode(buff_crop.getvalue()).decode("utf-8")

    #entire image for referring expression generation
    pil_I = Image.fromarray(I)
    buff = BytesIO()
    pil_I.save(buff, format='JPEG')
    base64_I = base64.b64encode(buff.getvalue()).decode("utf-8")

    # 구분 가능 여부 확인
    generator = OpenAI()
    response_check = generator.chat.completions.create(
        model="chatgpt-4o-latest",
        messages=[
            {
                "role": "user",
                "content": [
                    {

                        "type": "text",
                        "text": f"""Can the {cat_name} in the provided cropped image be clearly identified as belonging to the category {cat_name}? 
                                    Focus on whether the cropped image provides enough visible features (e.g., ears, head shape, fur texture) to confirm that it is a {cat_name}, even if the full body is not visible.

                                    Guidelines:
                                    - If the visible features (like ears, fur texture or head shape) are sufficient to identify the {cat_name}, respond with YES.
                                    - If multiple {cat_name}s are entangled or overlapping, making it difficult to distinguish one from another, respond with NONE.
                                    - If the object is clearly visible and identifiable as a {cat_name}, respond with YES.

                                    Output only either YES or NONE.
                        """
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_cropped_I}"},
                    }
                ]
            },
        ]
    )

    response_check_content = response_check.choices[0].message.content.strip().lower()
    #print(f"is object {obj_id} visible: {response_check_content}")
    
    if "yes" not in response_check_content:
        print(f"Referring expression not generated: {cat_name} is ambiguous in this frame.")
        return {"ref_exp": "NONE", "caption": caption, "cat_name": cat_name, "file_name": frame_name, "isValid" : False}

    # Referring expression 만들기
    # generator = OpenAI()
    response = generator.chat.completions.create(
        model="chatgpt-4o-latest",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",

                        "text": f"""Based on the dense caption, create a referring expression for the {cat_name} highlighted with the red box, corresponding to Object ID {obj_id}. 
                        Guidelines for creating the referring expression:
                        1. The referring expression should describe the prominent actions or poses of the highlighted {cat_name} (Object ID {obj_id}). 
                        2. Focus on the behavior or pose described in the caption that is specifically associated with this {cat_name}. Do not include actions or poses of other {cat_name}s.
                        3. If multiple {cat_name}s are present, ensure that the referring expression exclusively describes the {cat_name} corresponding to Object ID {obj_id}.
                        4. Avoid ambiguous or subjective terms. Use specific and clear action verbs to describe the highlighted {cat_name}.
                        5. The referring expression should only describe Object ID {obj_id} and not any other objects or entities. 
                        6. Use '{cat_name}' as the noun for the referring expressions. 
                        Output only the referring expression for the highlighted {cat_name} (Object ID {obj_id}).
                        
                        {caption}
                        """
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_I}"},
                    },
                    # {
                    #     "type": "image_url",
                    #     "image_url": {"url": f"data:image/jpeg;base64,{base64_cropped_I}"},
                    # }
                ],
            }
        ],
    )

    ref_exp = response.choices[0].message.content.strip()

    #QA filtering
    #QA1: 원하는 물체를 설명하는지
    filter = OpenAI()
    response1 = filter.chat.completions.create(
        model="chatgpt-4o-latest",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"""Does the given expression describe the {cat_name} highlighted with the red box? If so, only return YES and if not, NO.
                                    {ref_exp}""",
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_I}"},
                    },
                ],
            }
        ],
    )

    response1_content = response1.choices[0].message.content
    describesHighlighted = True if "yes" in response1_content.lower() else False

    #QA2: 원하지 않는 물체를 설명하지 않는지
    response2 = filter.chat.completions.create(
        model="chatgpt-4o-latest",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"""Does the given expression describe the person not highlighted with the red box? If so, only return YES and if not, NO.
                                    {ref_exp}""",
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_I}"},
                    },
                ],
            }
        ],
    )

    response2_content = response2.choices[0].message.content
    notDescribesNotHighlighted = False if "yes" in response2_content.lower() else True

    isValid = True if describesHighlighted and notDescribesNotHighlighted else False

    #print(f"describesHighlighted: {describesHighlighted}, notDescribesNotHighlighted: {notDescribesNotHighlighted}")
    #print(f"ref exp: {ref_exp}")
    #print("")

    return {"ref_exp": ref_exp, "caption": caption, "cat_name": cat_name, "file_name": frame_name, "isValid" : isValid}


if __name__ == '__main__':
    with open('mbench/sampled_frame3.json', 'r') as file:
        data = json.load(file)

    vid_ids = list(data.keys())
    all_ref_exps = {}

    os.environ['OPENAI_API_KEY'] = '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'

    # 전체 데이터셋의 vid_id에 대해
    for i in range(50):
        vid_id = vid_ids[i]
        
        #====캡션 만들기====
        # print("=====================captioner========================")
        captions, valid_obj_ids = getCaption(vid_id, data)
        cats_in_vid = list(captions.keys())
        # print()
        
        #====referring expression 만들고 QA filtering====
        # print("=====================referring expression generator & QA filter========================")
        ref_expressions = {}

        # 각 카테고리별로
        for cat_name in cats_in_vid:
            if cat_name not in ref_expressions:
                ref_expressions[cat_name] = {} 
            # 각 비디오 프레임 별로
            for frame_name in data[vid_id]['frame_names']:
                # print(f'--------category: {cat_name}, frame_name: {frame_name}')
                        
                if frame_name not in ref_expressions[cat_name]:
                    ref_expressions[cat_name][frame_name] = {}  # Create frame-level dictionary
                caption = captions[cat_name][frame_name]
                if not caption : continue
                else :
                    # 각 obj id별로
                    for obj_id in valid_obj_ids:
                        ref_exp = getRefExp(vid_id, frame_name, caption, obj_id, data)
                        ref_expressions[cat_name][frame_name][obj_id] = ref_exp  # Store ref_exp
                
        all_ref_exps[vid_id] = ref_expressions


    with open('mbench/result_revised50.json', 'w') as file:
        json.dump(all_ref_exps, file, indent=4)