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 | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from ultralytics import YOLO | |
| from src.tiles import * | |
| from src.image_transform import * | |
| from time import sleep | |
| ##### | |
| ### | |
| # 01 Setup | |
| ### | |
| ##### | |
| # This loads general code for use by the tool | |
| st.set_page_config( | |
| page_title="Find Waldo and His Friends Using Computer Vision", | |
| 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 'confidence' not in st.session_state: | |
| st.session_state.confidence = 0.50 | |
| 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, row5_col4 = st.columns([1,4,4,1], gap="medium") | |
| # heading | |
| with row0_col1: | |
| st.markdown("<h1 style='text-align: center;'>Find Waldo and His Friends</h1>", unsafe_allow_html=True) | |
| st.markdown("<h3 style='text-align: center;'>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;'>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 finds Waldo and his friends Wenda, Odlaw, Wizard, and Woof.</p>", unsafe_allow_html=True) | |
| # purchase option | |
| st.markdown("<p style='text-align: center;'><a href='https://www.amazon.com/Wheres-Waldo-Ultimate-Watcher-Collection/dp/1536215112/ref=asc_df_1536215112/?tag=hyprod-20&linkCode=df0&hvadid=496186854683&hvpos=&hvnetw=g&hvrand=13774168920524356768&hvpone=&hvptwo=&hvqmt=&hvdev=c&hvdvcmdl=&hvlocint=&hvlocphy=9009670&hvtargid=pla-1185030597547&psc=1&mcid=a517c47662b13f96a80fbf3bfe8b30fe&gclid=CjwKCAjwuJ2xBhA3EiwAMVjkVADBaOb87R0kxhnpKqL-S_gJWme0DJJAwSTe-VlUv2p19zZKak3cshoCR34QAvD_BwE'>Purchase the Books</a></p>", unsafe_allow_html=True) | |
| st.markdown("<p style='text-align: center;'><a href=mailto:'Jared.L.Bailey@duke.edu'>Jared.L.Bailey@duke.edu</a></p>", unsafe_allow_html=True) | |
| # proceed button | |
| _, row1a_col1, _ = st.columns([2,1,2]) | |
| with row1a_col1: | |
| if st.button("Proceed", key="Proceed", use_container_width=True): | |
| 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;'>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 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| # image selector | |
| st.markdown("<h4 style='text-align: center;'>Or Select One of the Following</h4>", unsafe_allow_html=True) | |
| image_path = "/home/user/app/test_images/" | |
| with row3_col1: | |
| img_1 = image_transform("IMG_5356.JPG", image_path) | |
| st.image(img_1, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_1", use_container_width=True): | |
| st.session_state.image = img_1 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row3_col2: | |
| img_2 = image_transform("IMG_5357.JPG", image_path) | |
| st.image(img_2, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_2", use_container_width=True): | |
| st.session_state.image = img_2 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row3_col3: | |
| img_3 = image_transform("IMG_5368.JPG", image_path) | |
| st.image(img_3, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_3", use_container_width=True): | |
| st.session_state.image = img_3 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row3_col4: | |
| img_4 = image_transform("IMG_5369.JPG", image_path) | |
| st.image(img_4, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_4", use_container_width=True): | |
| st.session_state.image = img_4 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col1: | |
| img_5 = image_transform("IMG_5382.JPG", image_path) | |
| st.image(img_5, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_5", use_container_width=True): | |
| st.session_state.image = img_5 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col2: | |
| img_6 = image_transform("IMG_5383.JPG", image_path) | |
| st.image(img_6, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_6", use_container_width=True): | |
| st.session_state.image = img_6 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col3: | |
| img_7 = image_transform("IMG_5408.JPG", image_path) | |
| st.image(img_7, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_7", use_container_width=True): | |
| st.session_state.image = img_7 | |
| st.session_state.screen = 2 | |
| st.rerun() | |
| with row4_col4: | |
| img_8 = image_transform("IMG_5409.JPG", image_path) | |
| st.image(img_8, use_column_width=True) | |
| if st.button("Select Image ↑", key="button_8", use_container_width=True): | |
| 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.session_state.boxed_image = np.array(st.session_state.image) | |
| height, width, channels = st.session_state.boxed_image.shape | |
| st.session_state.blank_img = np.zeros((height + 640, width + 640, 3), np.uint8) | |
| st.session_state.blank_img[0:height, 0:width] = st.session_state.boxed_image | |
| counter = 1 | |
| for i in range(0, width, 600): | |
| for j in range(0, height, 600): | |
| # draw the rectangle | |
| cv2.rectangle(st.session_state.blank_img, (i, j), (i + 640, j + 640), (0, 255, 0), 10) | |
| # add the number of the rectangle | |
| cv2.putText(st.session_state.blank_img, str(str(counter)), (i + 200, j + 320), cv2.FONT_HERSHEY_SIMPLEX, 10, (0, 255, 0), 10) | |
| counter += 1 | |
| st.image(st.session_state.blank_img, use_column_width=True) | |
| with row5_col3: | |
| # navigation | |
| row5a_col0, row5a_col1, row5a_col2= st.columns([1,1,1], gap="large") | |
| with row5a_col0: | |
| if st.button("Back", use_container_width=True): | |
| if st.session_state.image_counter > 0: | |
| st.session_state.image_counter -= 1 | |
| with row5a_col1: | |
| if st.button("Select New Photo", use_container_width=True): | |
| st.session_state.screen = 1 | |
| st.session_state.results = None | |
| st.session_state.image_counter = 0 | |
| st.rerun() | |
| with row5a_col2: | |
| if st.button("Next", use_container_width=True): | |
| 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 or st.session_state.results == None: | |
| st.write("\n\n") | |
| st.markdown("<p style='text-align: center;'>The model is working. Please be patient...</p>", unsafe_allow_html=True) | |
| st.markdown("<p style='text-align: center;'>This process can take up to 30 seconds.</p>", unsafe_allow_html=True) | |
| st.session_state.results = st.session_state.model.predict(st.session_state.tiles, conf=st.session_state.confidence) # predict | |
| st.rerun() | |
| with row5_col4: | |
| # show character locations | |
| class_names = {0: "Waldo", 1: "Wenda", 2: "Odlaw", 3: "Wizard", 4: "Woof"} | |
| character_dict = { | |
| "Waldo": [], | |
| "Wenda": [], | |
| "Odlaw": [], | |
| "Wizard": [], | |
| "Woof": [] | |
| } | |
| for idx, pred in enumerate(st.session_state.results): | |
| try: | |
| for boxes in pred.boxes: | |
| if boxes.conf >= 0.5: | |
| character_dict[class_names[int(boxes.cls.item())]].append(idx + 1) | |
| except: | |
| pass | |
| st.session_state.output = " \n\n\n\n\n Tile Numbers (There may be more than one example of a character in each image): \n" | |
| for key, value in character_dict.items(): | |
| if value != []: | |
| st.session_state.output += str(key) + ": "+ str(value) + " \n" | |
| st.write(st.session_state.output) | |
| # plot predicted image | |
| st.markdown(f"<h4 style='text-align: center;'>Image Tile {str(st.session_state.image_counter + 1)}</h4>", unsafe_allow_html=True) | |
| 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) |