Spaces:
Sleeping
Sleeping
| ##### | |
| ### | |
| # 00 Imports | |
| # 01 Setup | |
| # 02 Screen 0 - Introduction | |
| # 03 Screen 1 - Upload or Select Image | |
| # 04 Screen 2 - Output and Interaction | |
| ### | |
| ##### | |
| ##### | |
| ### | |
| # 00 Imports | |
| ### | |
| ##### | |
| import streamlit as st | |
| import streamlit.components.v1 as components | |
| from PIL import ImageOps, Image | |
| from ultralytics import YOLO | |
| from src.tiles import * | |
| ##### | |
| ### | |
| # 01 Setup | |
| ### | |
| ##### | |
| # This loads general code for use by the tool | |
| st.set_page_config(layout="wide") | |
| if 'screen' not in st.session_state: | |
| st.session_state.screen = 0 | |
| if 'image' not in st.session_state: | |
| st.session_state.image = None | |
| if 'image_counter' not in st.session_state: | |
| st.session_state.image_counter = 0 | |
| if 'model' not in st.session_state: | |
| st.session_state.model = YOLO("/home/user/app/models/head_model/best.pt") # load | |
| _, row0_col1, _ = st.columns([2,3,2]) | |
| _, row1_col1, _ = st.columns([2,3,2]) | |
| _, row2_col1, _ = st.columns([2,3,2]) | |
| _, row3_col1, row3_col2, row3_col3, row3_col4, _ = st.columns([8,3,3,3,3,8]) | |
| _, row4_col1, row4_col2, row4_col3, row4_col4, _ = st.columns([8,3,3,3,3,8]) | |
| _, row5_col2, row5_col3, _ = st.columns([1,4,4,1], gap="medium") | |
| # heading | |
| with row0_col1: | |
| st.markdown("<h1 style='text-align: center; color: black;'>Find Waldo and His Friends</h1>", unsafe_allow_html=True) | |
| st.markdown("<h3 style='text-align: center; color: black;'>Using Computer Vision</h3>", unsafe_allow_html=True) | |
| ##### | |
| ### | |
| # 02 Screen 0 - Introduction | |
| ### | |
| ##### | |
| # This displays to inform the user about the tool | |
| if st.session_state.screen == 0: | |
| with row1_col1: | |
| # overview | |
| st.markdown("<p style='text-align: left; color: black;'>This tool allows you to take a photo with your phone of a <i>Where's Waldo?</i> book page. Using the computer vision model YOLOv8-large, this tool find Waldo and his friends Wenda, Odlaw, Wizard, and Woof.</p>", unsafe_allow_html=True) | |
| # proceed button | |
| if st.button(f"Proceed", key="Proceed"): | |
| st.session_state.screen = 1 | |
| st.rerun() | |
| # image of Waldo and his friends | |
| st.image("/home/user/app/home_page_images/Waldo_Friends.jpg", use_column_width=True) | |
| ##### | |
| ### | |
| # 03 Screen 1 - Upload or Select Image | |
| ### | |
| ##### | |
| if st.session_state.screen == 1: | |
| # heading and instructions | |
| with row1_col1: | |
| st.markdown("<p style='text-align: center; color: black;'>Upload an image using your cell phone (single book page), or select an image from the below list</p>", unsafe_allow_html=True) | |
| # image uploader | |
| with row2_col1: | |
| uploaded_image = st.file_uploader(label="", | |
| type=['jpg'], | |
| accept_multiple_files=False, | |
| key=None, | |
| help='Due to image resolution limitations of the tool, only photos of a single page are expected to produce the intended result', | |
| on_change=None, | |
| args=None, | |
| kwargs=None, | |
| disabled=False, | |
| label_visibility="visible" | |
| ) | |
| if uploaded_image is not None: | |
| uploaded_image = Image.open(uploaded_image) | |
| uploaded_image = ImageOps.exif_transpose(uploaded_image) | |
| st.session_state.image = uploaded_image | |
| uploaded_image = None | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| # image selector | |
| st.markdown("<h4 style='text-align: center; color: black;'>Or Select One of the Following</h4>", unsafe_allow_html=True) | |
| image_path = "/home/user/app/test_images/" | |
| with row3_col1: | |
| img_1 = Image.open(image_path + "IMG_5356.JPG") | |
| img_1 = ImageOps.exif_transpose(img_1) | |
| st.image(img_1, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_1"): | |
| st.session_state.image = img_1 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row3_col2: | |
| img_2 = Image.open(image_path + "IMG_5357.JPG") | |
| img_2 = ImageOps.exif_transpose(img_2) | |
| st.image(img_2, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_2"): | |
| st.session_state.image = img_2 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row3_col3: | |
| img_3 = Image.open(image_path + "IMG_5368.JPG") | |
| img_3 = ImageOps.exif_transpose(img_3) | |
| st.image(img_3, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_3"): | |
| st.session_state.image = img_3 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row3_col4: | |
| img_4 = Image.open(image_path + "IMG_5369.JPG") | |
| img_4 = ImageOps.exif_transpose(img_4) | |
| st.image(img_4, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_4"): | |
| st.session_state.image = img_4 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col1: | |
| img_5 = Image.open(image_path + "IMG_5382.JPG") | |
| img_5 = ImageOps.exif_transpose(img_5) | |
| st.image(img_5, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_5"): | |
| st.session_state.image = img_5 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col2: | |
| img_6 = Image.open(image_path + "IMG_5383.JPG") | |
| img_6 = ImageOps.exif_transpose(img_6) | |
| st.image(img_6, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_6"): | |
| st.session_state.image = img_6 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col3: | |
| img_7 = Image.open(image_path + "IMG_5408.JPG") | |
| img_7 = ImageOps.exif_transpose(img_7) | |
| st.image(img_7, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_7"): | |
| st.session_state.image = img_7 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col4: | |
| img_8 = Image.open(image_path + "IMG_5409.JPG") | |
| img_8 = ImageOps.exif_transpose(img_8) | |
| st.image(img_8, use_column_width=True) | |
| if st.button(f"Select Image ↑", key="button_8"): | |
| st.session_state.image = img_8 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| ##### | |
| ### | |
| # 04 Screen 2 - Output and Interaction | |
| ### | |
| ##### | |
| if st.session_state.screen == 2: | |
| # tile images | |
| st.session_state.tiles = tile_image(image=st.session_state.image, tile_size=640, overlap=40) | |
| with row5_col2: | |
| st.image(st.session_state.image, use_column_width=True) | |
| with row5_col3: | |
| # navigation | |
| row5a_col0, row5a_col1, row5a_col2= st.columns([1,1,1], gap="medium") | |
| with row5a_col0: | |
| if st.button("Back"): | |
| if st.session_state.image_counter > 0: | |
| st.session_state.image_counter -= 1 | |
| with row5a_col1: | |
| if st.button("Select New Photo"): | |
| st.session_state.screen = 1 | |
| st.session_state.results = None | |
| st.rerun() | |
| with row5a_col2: | |
| if st.button("Next"): | |
| if st.session_state.image_counter < len(st.session_state.results) - 1: | |
| st.session_state.image_counter += 1 | |
| # predictions | |
| if 'results' not in st.session_state: | |
| st.write("The model is working. Please be patient...") | |
| st.write("This process takes around 30 seconds.") | |
| st.session_state.results = st.session_state.model.predict(st.session_state.tiles, conf=0.5) # predict | |
| st.rerun() | |
| # plot predicted image | |
| im_bgr = st.session_state.results[st.session_state.image_counter].plot() # BGR-order numpy array | |
| im_rgb = Image.fromarray(im_bgr[..., ::-1]) # RGB-order PIL image | |
| st.image(im_rgb) | |