captioning_gpt.py

# ==============================================================================
# Copyright (c) 2024 Tiange Luo, tiange.cs@gmail.com
# Last modified: September 04, 2024
#
# This code is licensed under the MIT License.
# ==============================================================================
import base64
import requests
import openai
import os
import pickle
import numpy as np
import csv
import pandas as pd
import glob
import argparse
from IPython import embed

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

parser = argparse.ArgumentParser(description="Process API key and CSV file path.")
parser.add_argument('--api_key', type=str, required=True, help="Your OpenAI API Key.")
parser.add_argument('--csv_file', type=str, default='./caption.csv', help="Path to the output CSV file.")
parser.add_argument("--parent_dir", type = str, default='./example_material/Cap3D_imgs')
args = parser.parse_args()

api_key = args.api_key
csv_file = args.csv_file

output_csv = open(csv_file, 'a')
writer = csv.writer(output_csv)

paths = glob.glob(os.path.join(args.parent_dir, '*'))
wrong_or_none_files = []
captions = {}
for index, path in enumerate(paths):
    uid = path.split('/')[-1]
    image_paths = []
    # insert your image_path
    # in DiffuRank, we send the top-6 views after ranking with DiffuRank
    diffurank_scores = pickle.load(open(os.path.join(path, 'diffurank_scores.pkl'), 'rb'))
    ranks = np.argsort(diffurank_scores)

    for i in range(6):
        image_paths.append(os.path.join(path, '%05d.png'%ranks[i]))

    base64_images = []
    try:
        for p in image_paths:
            base64_images.append(encode_image(p))
    except:
        wrong_or_none_files.append(u)
        print('image enconding error')
        continue
    
    headers = {
      "Content-Type": "application/json",
      "Authorization": f"Bearer {api_key}"
    }
    
    payload = {
      # GPT-4o-mini is much cheaper than GPT-4o, 
      # DiffuRank captions are generated by GPT-4o
      # "model": "gpt-4o-2024-05-13",
      "model": "gpt-4o-mini-2024-07-18",
      "messages": [
        {
          "role": "user",
          "content": [
            {
              "type": "text",
              "text": "Renderings show different angles of the same set of 3D objects. Concisely describe 3D object (distinct features, objects, structures, material, color, etc) as a caption, not mentioning angles and image related words"
            },
          ]
        }
      ],
      "max_tokens": 300
    }
    for i in range(len(base64_images)):
        payload['messages'][0]['content'].append( {
       "type": "image_url",
       "image_url": {
         "url": f"data:image/jpeg;base64,{base64_images[i]}"
       }
     }
    )
    
    try:
        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
    except:
        continue
    try:
        r = response.json()
    except:
        continue
    captions[uid] = r
    try:
        cur_caption = r['choices'][0]['message']['content']
    except:
        continue
    writer.writerow([uid, cur_caption])
    print(index, uid, cur_caption)
    if (index)% 100 == 0:
        output_csv.flush()
        os.fsync(output_csv.fileno())
output_csv.close()