Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from crowd_counter import CrowdCounter | |
| from PIL import Image | |
| import cv2 | |
| import numpy as np | |
| # Initialize the crowd counter engine | |
| crowd_counter_engine = CrowdCounter() | |
| # Function to predict crowd count and generate prediction image | |
| def predict(input_image): | |
| input_image = Image.fromarray(input_image.astype('uint8'), 'RGB') | |
| response = crowd_counter_engine.inference(input_image) | |
| crowd_count, pred_img = response | |
| pred_img_rgb = cv2.cvtColor(pred_img, cv2.COLOR_BGR2RGB) | |
| return pred_img_rgb, crowd_count | |
| # Streamlit UI layout | |
| st.markdown("<h1 style='text-align: center; margin-bottom: 0.5rem;'>Crowd Counter Demo</h1>", unsafe_allow_html=True) | |
| st.markdown("<h3 style='text-align: center; margin-bottom: 1rem;'>A Demo of Proposal Point Prediction for Crowd Counting</h3>", unsafe_allow_html=True) | |
| # Radio button to choose between example images and uploading an image | |
| option = st.radio('Choose an image source:', ('Use Example Image', 'Upload Your Own Image')) | |
| # Example images | |
| example_images = { | |
| "Example 1": "images/img-1.jpg", | |
| "Example 2": "images/img-2.jpg", | |
| "Example 3": "images/img-3.jpg", | |
| } | |
| # Handling the selection | |
| if option == 'Use Example Image': | |
| # Select box for example images | |
| example_selection = st.selectbox('Choose an example image:', list(example_images.keys())) | |
| if example_selection: | |
| example_image_path = example_images[example_selection] | |
| example_image = Image.open(example_image_path) | |
| st.image(example_image, caption='Selected Example Image', use_column_width=True) | |
| # Convert the PIL Image to an array for processing | |
| example_image_array = np.array(example_image) | |
| with st.spinner("Thinking..."): | |
| # Predict the crowd count and proposal points for the example image | |
| pred_img, crowd_count = predict(example_image_array) | |
| # Display the prediction results | |
| st.success(f"Predicted Count: {crowd_count}") | |
| st.image(pred_img, caption='Proposal Points Prediction', use_column_width=True) | |
| else: # 'Upload Your Own Image' is selected | |
| uploaded_file = st.file_uploader("Or upload an image:", type=['jpg', 'jpeg', 'png']) | |
| if uploaded_file is not None: | |
| file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8) | |
| user_img = cv2.imdecode(file_bytes, 1) | |
| st.image(user_img, caption='Uploaded Image', use_column_width=True) | |
| with st.spinner("Thinking..."): | |
| # Predict the crowd count and proposal points for the uploaded image | |
| pred_img, crowd_count = predict(user_img) | |
| # Display the prediction results | |
| st.success(f"Predicted Count: {crowd_count}") | |
| st.image(pred_img, caption='Proposal Points Prediction', use_column_width=True) | |
| st.markdown('---') | |
| footer = """ | |
| <style> | |
| .footer { | |
| font-family: Arial, sans-serif; | |
| font-size: small; | |
| text-align: center; | |
| padding: 10px; | |
| margin-top: 20px; | |
| } | |
| </style> | |
| <div class="footer"> | |
| © 2023 <a href="https://www.ristek.cs.ui.ac.id/" target="_blank">RISTEK Fasilkom UI</a> <br>All rights reserved<br> | |
| Powered by <a href="https://github.com/TencentYoutuResearch/CrowdCounting-P2PNet/" target="_blank">P2PNet</a> | |
| </div> | |
| """ | |
| st.markdown(footer, unsafe_allow_html=True) |