-
Notifications
You must be signed in to change notification settings - Fork 1
/
app.py
57 lines (47 loc) · 1.75 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import streamlit as st
from PIL import Image
from generate_DSL import generate_code
from utils import *
from generate_DSL import *
from generate_HTML import *
import textwrap
def to_markdown(text):
text = text.replace('•', ' *')
return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))
# Set page title
st.set_page_config(page_title="Sketch to UI Converter", layout="wide")
# Title
st.title("Sketch to UI Converter")
# Sidebar
st.sidebar.header("Settings")
uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
# Load the tokenizer, ViT model, and decoder
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.add_special_tokens(special_tokens_dict)
vit_model = ViTModel.from_pretrained(
'google/vit-base-patch16-224').base_model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vit_model.to(device)
decoder = GPT2DecoderWithImageFeatures(input_size=768)
# Update the GPT2 model with the new tokenizer
decoder.gpt.resize_token_embeddings(len(tokenizer))
decoder.load_state_dict(torch.load(
"best_decoder.pth", map_location=device))
decoder.to(device)
# Main content
while bool(uploaded_file)==True:
# Display uploaded image
st.image(uploaded_file, caption="Uploaded Image.", use_column_width=True)
img_rgb = Image.open(uploaded_file)
# Perform Conversion
code_result = generate_code(img_rgb, tokenizer, vit_model, decoder)
dsl_code = '''{code_result}'''
# Display extracted text
st.subheader("Generated DSL Code:")
st.code(code_result, language='plaintext')
# Display HTML/CSS code
st.subheader("Generated HTML/CSS Code:")
Gemini_HTML_code = generate_HTML(code_result)
HTML_code = Gemini_HTML_code.text
st.markdown(HTML_code)
break