Update app.py

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import re
 import os 
 import json
+from typing import Union
 hf_token = os.environ.get('HF_TOKEN')
 
 from gradio_client import Client, handle_file
@@ -181,29 +182,32 @@ def parse_perfume_description(text: str) -> dict:
 
     return result
 
-def extract_image_description(data: dict) -> str:
+def extract_field(data: Union[str, dict], field_name: str) -> str:
     """
-    Extracts the Image Description value from a Python dict response.
+    Extracts a specific field value from a JSON string or Python dict.
 
     Args:
-        data (dict): The dict returned by your LLM.
+        data (Union[str, dict]): The JSON string or dict to extract from.
+        field_name (str): The exact field name to extract.
 
     Returns:
-        str: The Image Description text.
+        str: The extracted field value as a string.
     """
+    if isinstance(data, str):
+        try:
+            data = json.loads(data)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON string provided")
+
     if not isinstance(data, dict):
-        raise TypeError("Input must be a dict")
+        raise TypeError("Input must be a dict or a valid JSON string")
 
-    image_description = (
-        data.get("Image Description")
-        or data.get("image_description")
-        or None
-    )
+    value = data.get(field_name) or data.get(field_name.lower()) or None
 
-    if not image_description:
-        raise KeyError("No 'Image Description' field found in the response")
+    if value is None:
+        raise KeyError(f"No field named '{field_name}' found in the data")
 
-    return image_description.strip()
+    return str(value).strip()
 
 
 def get_text_after_colon(input_text):
@@ -240,7 +244,8 @@ def infer(image_input):
 
     parsed = parse_perfume_description(result)
 
-    image_desc = extract_image_description(parsed)
+    image_desc = extract_field(parsed, "Image Description")
+
     print(image_desc)
 
     return result, parsed