try streamlit

crowd_counter/__init__.py

@@ -69,19 +69,15 @@ class CrowdCounter:
     ) -> tuple[any, Image.Image, torch.Tensor]:
 
         ori_width, ori_height = img_raw.size
-        max_dimension = 384
-
-        if ori_width > max_dimension or ori_height > max_dimension:
-            scale_factor = max_dimension / max(ori_width, ori_height)
-            new_width = int(ori_width * scale_factor)
-            new_height = int(ori_height * scale_factor)
-            new_width = new_width // 128 * 128
-            new_height = new_height // 128 * 128
-            img_resized = img_raw.resize((new_width, new_height), Image.LANCZOS)
-        else:
-            img_resized = img_raw
-            new_width = ori_width
-            new_height = ori_height
+        max_dimension = 512
+
+        scale_factor = max_dimension / max(ori_width, ori_height)
+        new_width = int(ori_width * scale_factor)
+        new_height = int(ori_height * scale_factor)
+        print(new_width, new_height)
+        new_width = new_width // 128 * 128
+        new_height = new_height // 128 * 128
+        img_resized = img_raw.resize((new_width, new_height), Image.LANCZOS)
 
         print(new_width, new_height)
         # pre-proccessing

main.py

@@ -0,0 +1,88 @@
+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">
+            &copy; 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)

requirements.txt

@@ -1,11 +1,12 @@
-torch
-torchvision
-tensorboardX
-easydict
-pandas
-numpy
-scipy
-matplotlib
-Pillow
-opencv-python
-gradio==3.50
+torch
+torchvision
+tensorboardX
+easydict
+pandas
+numpy
+scipy
+matplotlib
+Pillow
+opencv-python
+gradio==3.50
+# streamlit