Svngoku commited on
Commit
97f4a42
·
verified ·
1 Parent(s): e4fb89e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +65 -0
app.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from smolagents.agents import ToolCallingAgent
3
+ from smolagents import tool, LiteLLMModel
4
+ from typing import Optional
5
+ import cv2
6
+ import pytesseract
7
+ from PIL import Image
8
+ import io
9
+ import numpy as np
10
+ import base64
11
+
12
+ # Define the LiteLLMModel with OpenAI key
13
+ model = LiteLLMModel(model_id="gpt-4o", api_key="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")
14
+
15
+ @tool
16
+ def extract_components(image_data_base64: str) -> str:
17
+ """
18
+ Extract components from a web design image.
19
+
20
+ Args:
21
+ image_data_base64: The image data in base64 string format.
22
+
23
+ Returns:
24
+ A string describing the components found in the image.
25
+ """
26
+ image_data = base64.b64decode(image_data_base64)
27
+ image = Image.open(io.BytesIO(image_data))
28
+ gray = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)
29
+ components = pytesseract.image_to_string(gray)
30
+ return components
31
+
32
+ @tool
33
+ def generate_code(components: str) -> str:
34
+ """
35
+ Generate code for the given components.
36
+
37
+ Args:
38
+ components: A string describing the components.
39
+
40
+ Returns:
41
+ The generated code for the components.
42
+ """
43
+ # This is a placeholder implementation. You can replace it with actual code generation logic.
44
+ return f"Generated code for components: {components}"
45
+
46
+ # Define the ToolCallingAgent
47
+ agent = ToolCallingAgent(tools=[extract_components, generate_code], model=model)
48
+
49
+ # Streamlit app title
50
+ st.title("Web Design Component Extractor")
51
+
52
+ # File uploader for the web design image
53
+ uploaded_file = st.file_uploader("Upload a web design image", type=["png", "jpg", "jpeg"])
54
+
55
+ # Button to run the agent
56
+ if st.button("Extract and Generate Code"):
57
+ if uploaded_file is not None:
58
+ image_data = uploaded_file.read()
59
+ image_data_base64 = base64.b64encode(image_data).decode('utf-8')
60
+ components = agent.run(f"extract_components {image_data_base64}")
61
+ code = agent.run(f"generate_code {components}")
62
+ st.write("Extracted Components:", components)
63
+ st.write("Generated Code:", code)
64
+ else:
65
+ st.write("Please upload an image.")