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("Crowd Counter Demo
", unsafe_allow_html=True)
st.markdown("A Demo of Proposal Point Prediction for Crowd Counting
", 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 = """
"""
st.markdown(footer, unsafe_allow_html=True)