.history/mbench/gpt_ref-ytvos_20250119070824.py

from datasets import build_dataset
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

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

# 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")

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()
    for obj_id in list(video_data['annotations'][0].keys()):
        cat_names.add(video_data['annotations'][0][obj_id]['category_name'])
    
    if len(cat_names) == 1:
        cat_name = next(iter(cat_names))
    else:
        print("more than 2 categories")
        return -1
    
    image_paths = [os.path.join(video_path, frame_name + '.jpg') for frame_name in frame_names]
    image_captions = {}

    captioner = OpenAI()
    for i in range(len(image_paths)):
        image_path = image_paths[i]
        frame_name = frame_names[i]
        base64_image = encode_image(image_path)
        
        #1단계: 필터링
        response1 = captioner.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": f"Are there multiple {cat_name}s that can be distinguished by action? Each action should be prominent and describe the corresponding object only. If so, only output YES. If not, only output 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

        #2단계: dense caption 만들기
        if should_caption:
            response2 = captioner.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": f"""
                                    Describe the image in detail focusing on the {cat_name}s' actions. 
                                    1. Each action should be prominent, clear and unique, describing the corresponding object only.
                                    2. Avoid overly detailed or indeterminate details such as ‘in anticipation’. 
                                    3. Avoid subjective descriptions such as ‘soft’, ‘controlled’, ‘attentive’, ‘skilled’, ‘casual atmosphere’ and descriptions of the setting. 
                                    4. Do not include actions that needs to be guessed or suggested.""",
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                            },
                        ],
                    }
                ],
            )

            caption = response2.choices[0].message.content
        else:
            caption = None

        image_captions[frame_name] = caption
    return image_captions

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()
    pil_I = Image.fromarray(I)
    buff = BytesIO()
    pil_I.save(buff, format='JPEG')
    base64_I = base64.b64encode(buff.getvalue()).decode("utf-8")

    #ref expression 만들기
    generator = OpenAI()
    response = generator.chat.completions.create(
        model="gpt-4o-mini",
        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. 
                                1. The referring expression describes the action and does not contain information about appearance or location in the picture. 
                                2. Focus only on prominent actions and avoid overly detailed or indeterminate details. 
                                3. Avoid subjective terms describing emotion such as ‘in anticipation’, ‘attentively’ or ‘relaxed’ and professional, difficult words. 
                                4. The referring expression should only describe the highlighted {cat_name} and not any other. 
                                5. Use '{cat_name}' as the noun for the referring expressions. 
                                Output only the referring expression.
                                {caption}""",
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_I}"},
                    },
                ],
            }
        ],
    )

    ref_exp = response.choices[0].message.content

    #QA filtering
    #QA1: 원하는 물체를 설명하는지
    filter = OpenAI()
    response1 = filter.chat.completions.create(
        model="gpt-4o-mini",
        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="gpt-4o-mini",
        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
    describesNotHighlighted = True if "yes" in response2_content.lower() else False

    isValid = True if describesHighlighted and not describesNotHighlighted else False

    print(f"describesHighlighted: {describesHighlighted}, describesNotHighlighted: {describesNotHighlighted}")

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

def createRefExp(video_id, json_data):
    video_data = json_data[video_id]
    obj_ids = list(video_data['annotations'][0].keys())
    frame_names = video_data['frame_names']

    captions_per_frame = getCaption(video_id, json_data)
    
    if captions_per_frame == -1:
        print("There are more than 2 cateories")
        return
    

    video_ref_exps = {}

    for frame_name in frame_names:
        frame_caption = captions_per_frame[frame_name]

        if frame_caption == None:
            video_ref_exps[frame_name] = None

        else:
            frame_ref_exps = {}
            for obj_id in obj_ids:
                exp_per_obj = getRefExp(video_id, frame_name, frame_caption, obj_id, json_data)
                frame_ref_exps[obj_id] = exp_per_obj
            video_ref_exps[frame_name] = frame_ref_exps

    return video_ref_exps     

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

    videos = set()
    with open('make_ref-ytvos/selected_frames.jsonl', 'r') as file:
        manual_select = list(file)
    for frame in manual_select:
        result = json.loads(frame)
        videos.add(result['video'])
    videos = list(videos)


    all_video_refs = {}
    for i in range(1):
        video_id = videos[i]
        video_ref = createRefExp(video_id, data)
        all_video_refs[video_id] = video_ref