Update app.py

app.py

@@ -7,21 +7,23 @@ import pytesseract
 from PIL import Image
 import io
 import numpy as np
+import base64
 
 # Define the LiteLLMModel with OpenAI key
 model = LiteLLMModel(model_id="gpt-4o", api_key="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")
 
 @tool
-def extract_components(image_data: bytes) -> str:
+def extract_components(image_data_base64: str) -> str:
     """
     Extract components from a web design image.
 
     Args:
-        image_data: The image data in bytes.
+        image_data_base64: The image data in base64 string format.
 
     Returns:
         A string describing the components found in the image.
     """
+    image_data = base64.b64decode(image_data_base64)
     image = Image.open(io.BytesIO(image_data))
     gray = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)
     components = pytesseract.image_to_string(gray)
@@ -54,8 +56,9 @@ uploaded_file = st.file_uploader("Upload a web design image", type=["png", "jpg"
 if st.button("Extract and Generate Code"):
     if uploaded_file is not None:
         image_data = uploaded_file.read()
-        components = agent.run("extract_components", image_data=image_data)
-        code = agent.run("generate_code", components=components)
+        image_data_base64 = base64.b64encode(image_data).decode('utf-8')
+        components = agent.run(f"extract_components {image_data_base64}")
+        code = agent.run(f"generate_code {components}")
         st.write("Extracted Components:", components)
         st.write("Generated Code:", code)
     else: