import sys from os import path as osp sys.path.append(osp.abspath(osp.join(osp.dirname(__file__), '..'))) from datasets import build_dataset import argparse import opts from pathlib import Path import os 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 None 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, 2): video_id = videos[i] video_ref = createRefExp(video_id, data) all_video_refs[video_id] = video_ref json_obj = json.dumps(all_video_refs, indent=4) with open('mbench/result.json', 'w') as file: file.write(json_obj)