134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
import os
|
|
|
|
import requests
|
|
import streamlit as st
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
from typing import List
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class MosaicRequest(BaseModel):
|
|
rows: int
|
|
columns: int
|
|
buffer_size: int
|
|
keywords: List[str]
|
|
background_file: str
|
|
|
|
|
|
class BackgroundRequest(BaseModel):
|
|
backgrounds: List[str]
|
|
|
|
|
|
backend_url = os.environ.get('BACKEND')
|
|
|
|
st.set_page_config(page_title='Mosaic Generator')
|
|
|
|
st.title('Mosaic Generator')
|
|
|
|
st.sidebar.title('Navigation')
|
|
tabs = st.sidebar.radio('Select task', ['Generate Mosaic', 'Upload Image', 'Upload Background'])
|
|
|
|
if tabs == 'Generate Mosaic':
|
|
st.header('Generate Mosaic')
|
|
rows = st.number_input('Rows', min_value=1, value=10)
|
|
columns = st.number_input('Columns', min_value=1, value=10)
|
|
buffer_size = st.number_input('Buffer Size', min_value=1, value=100)
|
|
keywords = st.text_input('Keywords (separated by commas)', 'cat,dog')
|
|
# Send a GET request to the endpoint
|
|
response = requests.get(backend_url + '/backgrounds')
|
|
|
|
# Check if the request was successful
|
|
if response.status_code == 200:
|
|
# Parse the response JSON into a BackgroundRequest object
|
|
background_request = BackgroundRequest.parse_raw(response.text)
|
|
|
|
# Create a Streamlit selectbox with the list of backgrounds
|
|
background_file = st.selectbox("Select a background", options=background_request.backgrounds)
|
|
else:
|
|
# Display an error message if the request was not successful
|
|
background_file = ''
|
|
st.error("Failed to retrieve background options.")
|
|
|
|
if st.button('Generate Mosaic'):
|
|
if background_file is None:
|
|
st.warning('Please select a background image')
|
|
else:
|
|
|
|
# Send request to generate_mosaic endpoint
|
|
data = {
|
|
'rows': rows,
|
|
'columns': columns,
|
|
'buffer_size': buffer_size,
|
|
'keywords': keywords.split(','),
|
|
'background_file': background_file
|
|
}
|
|
mosaic_request = MosaicRequest(**data)
|
|
response = requests.post(backend_url + '/generate_mosaic',
|
|
json=mosaic_request.dict(),
|
|
headers={"Content-Type": "application/json"}
|
|
)
|
|
|
|
if response.ok:
|
|
st.success('Mosaic generated')
|
|
st.json(response.json())
|
|
else:
|
|
st.error('Error generating mosaic: ' + response.reason)
|
|
|
|
elif tabs == 'Upload Image':
|
|
st.header('Upload Image')
|
|
uploaded_file = st.file_uploader('Image File', type=['jpg', 'jpeg', 'png'])
|
|
|
|
if st.button('Upload Image'):
|
|
if uploaded_file is None:
|
|
st.warning('Please upload an image file')
|
|
else:
|
|
# Convert uploaded image to bytes
|
|
img_bytes = BytesIO()
|
|
Image.open(uploaded_file).save(img_bytes, format='PNG')
|
|
img_bytes.seek(0)
|
|
|
|
# Access the filename of the uploaded file
|
|
filename = uploaded_file.name
|
|
|
|
# Change the extension of the filename to 'png'
|
|
filename = os.path.splitext(filename)[0] + '.png'
|
|
|
|
# Send request to upload_file endpoint
|
|
files = {'file': (filename, img_bytes, 'image/png')}
|
|
response = requests.post(backend_url + '/upload_file', files=files)
|
|
|
|
if response.ok:
|
|
st.success('Image file uploaded')
|
|
st.json(response.json())
|
|
else:
|
|
st.error('Error uploading image file')
|
|
|
|
elif tabs == 'Upload Background':
|
|
st.header('Upload Background')
|
|
uploaded_file = st.file_uploader('Background Image', type=['jpg', 'jpeg', 'png'])
|
|
|
|
if st.button('Upload Background'):
|
|
if uploaded_file is None:
|
|
st.warning('Please upload a background image')
|
|
else:
|
|
# Convert uploaded image to bytes
|
|
img_bytes = BytesIO()
|
|
Image.open(uploaded_file).save(img_bytes, format='PNG')
|
|
img_bytes.seek(0)
|
|
|
|
# Access the filename of the uploaded file
|
|
filename = uploaded_file.name
|
|
|
|
# Change the extension of the filename to 'png'
|
|
filename = os.path.splitext(filename)[0] + '.png'
|
|
|
|
# Send request to upload_file endpoint
|
|
files = {'file': (filename, img_bytes, 'image/png')}
|
|
response = requests.post(backend_url + '/upload_background', files=files)
|
|
|
|
if response.ok:
|
|
st.success('Background image uploaded')
|
|
st.json(response.json())
|
|
else:
|
|
st.error('Error uploading background image')
|