adding main captioning functionality

app.py

@@ -1,4 +1,69 @@
 import gradio as gr
+import os, sys
+from prompting import promptingutils
+from imageprocessing import imageprocessingtools 
+from openai import OpenAI
+from prompting.promptingutils import DEFAULT_N_SAMPLES, DEFAULT_OBJECT_THRESHOLD, DEFAULT_RANDOM_STATE
+
+
+AVAILABLE_LLMS = [
+    "vicuna-7b",
+    "llama-7b-chat",
+    "mistral-7b-instruct",
+    "vicuna-13b",
+]
+
+DEFAULT_TEMPERATURE = 0
+LLAMA_API_TOKEN = os.environ["LLAMA_API_TOKEN"]
+
+client = OpenAI(
+    api_key = LLAMA_API_TOKEN,
+    base_url = "https://api.llama-api.com"
+)
+
+
+def caption_artwork(
+        image_filepath: os.PathLike,
+        llm :str, 
+        temperature = DEFAULT_TEMPERATURE,
+        items_threshold = DEFAULT_OBJECT_THRESHOLD,
+        random_state = DEFAULT_RANDOM_STATE,
+        n_samples_per_emotion = DEFAULT_N_SAMPLES
+)-> tuple:
+    
+    all_information = imageprocessingtools.extract_all_information_from_image(image_filepath)
+    
+    emotion = all_information["emotion"]
+    colors_list = all_information["colors_list"]
+    objects_and_probs = all_information["objects_and_probs"]
+    objects_list = promptingutils.filter_items(objects_and_probs, items_threshold=items_threshold)
+
+    user_prompt = promptingutils.get_user_prompt(
+        colors_list=colors_list, 
+        objects_list=objects_list,
+        emotion=emotion,
+        n_samples_per_emotion=n_samples_per_emotion,
+        random_state=random_state,
+        object_threshold=items_threshold
+        
+        )
+    
+    response = client.chat.completions.create(
+        model = llm,
+        messages = [
+            {"role": "system"   , "content": "Assistant is a large language model trained by OpenAI."},
+            {"role": "user"     , "content": user_prompt}
+        ],
+        temperature = temperature
+    )
+
+    commentary_str = response.choices[0].message.content
+    colors_str = ", ".join(colors_list)
+    objects_str = ", ".join(objects_list)
+    emotion_str = emotion
+
+    return (emotion_str, colors_str, objects_str, commentary_str)
+
 
 def greet(name):
     return "Hello " + name + "!!"

imageprocessing/imageprocessingtools.py

@@ -51,7 +51,7 @@ def extract_all_information_from_image(
     )
 
     result = {
-        "colors": colors,
+        "colors_list": colors,
         "objects_and_probs" : objects_and_probs,
         "emotion": emotion
     }

prompting/promptingutils.py

@@ -53,15 +53,15 @@ def fill_extracted_items(
 def load_dataframe(
     csv_filepath,
 ):
-  df = pd.read_csv(csv_filepath, index_col = 0)
-  for stringified_col in  [
-                            "maskrcnn_objects",
-                            "colors",
-                            "clip_recognized_objects", 
-                          ]:
-      df[stringified_col] = df[stringified_col].apply(eval)
+    df = pd.read_csv(csv_filepath, index_col = 0)
+    for stringified_col in  [
+                                "maskrcnn_objects",
+                                "colors",
+                                "clip_recognized_objects", 
+                            ]:
+        df[stringified_col] = df[stringified_col].apply(eval)
 
-  return df
+    return df
 
 TOP_COMMENTARIES_DFS = {
     emotion : load_dataframe(os.path.join(TOP_COMMENTARIES_DIR, f"top_{emotion.replace(' ', '_')}.csv"))
@@ -111,18 +111,18 @@ def get_subprompt_for_emotion(
     random_state = DEFAULT_RANDOM_STATE,
     object_threshold = DEFAULT_OBJECT_THRESHOLD,
 ):
-  random_samples = get_random_samples_for_emotion(
-      emotion = emotion,
-      n_samples = n_samples,
-      random_state = random_state,
-      object_threshold=object_threshold,
-  )
-  subprompt = [
-        fill_extracted_items(**entry) for entry in random_samples    
-  ]
-  subprompt = "\n".join(subprompt)
+    random_samples = get_random_samples_for_emotion(
+        emotion = emotion,
+        n_samples = n_samples,
+        random_state = random_state,
+        object_threshold=object_threshold,
+    )
+    subprompt = [
+            fill_extracted_items(**entry) for entry in random_samples    
+    ]
+    subprompt = "\n".join(subprompt)
 
-  return subprompt
+    return subprompt
 
 
 def get_subprompt_with_examples(
@@ -153,38 +153,38 @@ def get_user_prompt(
     random_state = DEFAULT_RANDOM_STATE,
     object_threshold = DEFAULT_OBJECT_THRESHOLD,
 ):
-  user_prompt= (
-      "You have to write a commentary for an artwork.\n"
-      "To write the commentary, you are given the objects present in the picture, "
-      "the colors present in the picture, and the emotion the picture evokes.\n"
-      "You are first shown several examples, and then have to give your commentary.\n"
-      "First come the examples, and then the objects, colors, and emotion you will have to use for your commentary.\n"
-      "Avoid explicitly mentioning the objects, or colors, or emotion, if it sounds more natural.\n"
-      "Only write the commentary.\n"
-      "\n"
-      "EXAMPLES:"
-      "\n\n"
-      "{examples}"
-      "\n"
-      "Now, write your personal opinion about the picture."
-      "\n"
-      "{image_subprompt}"
-  )
-
-  examples = get_subprompt_with_examples(
-      n_samples_per_emotion = n_samples_per_emotion,  
-      random_state = random_state,
-      object_threshold=object_threshold,
-  )
-
-  image_subprompt = fill_extracted_items(
-    colors_list = colors_list,
-    objects_list = objects_list,
-    emotion = emotion,
-    commentary = None,
-  )
-
-
-  result = user_prompt.format(examples = examples, image_subprompt = image_subprompt)
-
-  return result
+    user_prompt= (
+        "You have to write a commentary for an artwork.\n"
+        "To write the commentary, you are given the objects present in the picture, "
+        "the colors present in the picture, and the emotion the picture evokes.\n"
+        "You are first shown several examples, and then have to give your commentary.\n"
+        "First come the examples, and then the objects, colors, and emotion you will have to use for your commentary.\n"
+        "Avoid explicitly mentioning the objects, or colors, or emotion, if it sounds more natural.\n"
+        "Only write the commentary.\n"
+        "\n"
+        "EXAMPLES:"
+        "\n\n"
+        "{examples}"
+        "\n"
+        "Now, write your personal opinion about the picture."
+        "\n"
+        "{image_subprompt}"
+    )
+
+    examples = get_subprompt_with_examples(
+        n_samples_per_emotion = n_samples_per_emotion,  
+        random_state = random_state,
+        object_threshold=object_threshold,
+    )
+
+    image_subprompt = fill_extracted_items(
+        colors_list = colors_list,
+        objects_list = objects_list,
+        emotion = emotion,
+        commentary = None,
+    )
+
+
+    result = user_prompt.format(examples = examples, image_subprompt = image_subprompt)
+
+    return result