Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,6 +3,8 @@ from PIL import Image
 import time
 from ultralytics import YOLO
 import io
+import matplotlib.pyplot as plt
+import pandas as pd
 
 # --- Page Config ---
 st.set_page_config(page_title="Smart Traffic Light System", layout="wide")
@@ -27,7 +29,7 @@ uploaded_images = {}
 for col, direction in zip(cols, directions):
     with col:
         uploaded_images[direction] = st.file_uploader(f"{direction}", type=["jpg", "png"], key=direction)
-
+        
 # --- Process Once All Images Are Uploaded ---
 if all(uploaded_images.values()):
     # Initialize session state attributes if not already initialized
@@ -64,6 +66,7 @@ if all(uploaded_images.values()):
     yellow_time = 3  # Fixed yellow duration
 
     # --- Enhanced Signal Status Visualization ---
+    st.markdown("<br>", unsafe_allow_html=True)
     st.markdown("### 🚥 Current Signal Status")
 
     signal_cols = st.columns(4)
@@ -102,10 +105,12 @@ if all(uploaded_images.values()):
                 st.markdown(f"### {status}\n**{direction}**\n**{count} Vehicles**")
 
     # --- Timer UI ---
+    st.markdown("<br>", unsafe_allow_html=True)
     st.markdown("### ⏱ Signal Timer")
     timer_placeholder = st.empty()
 
     # --- Display Annotated Images Below Timer ---
+    st.markdown("<br>", unsafe_allow_html=True)
     st.markdown("### 🖼️ Detected Vehicles (Annotated Images)")
     img_cols = st.columns(4)
     for idx, direction in enumerate(directions):
@@ -138,59 +143,136 @@ if all(uploaded_images.values()):
         else:
             # Show completion message when all directions are done
             st.success("### 🚦 Simulation Complete! All directions have completed their green and yellow phases.")
-            # Add restart button
+            counts = st.session_state.counts
+            sorted_directions = st.session_state.sorted_directions
+            # st.session_state.clear()
+            
             st.markdown("<br>", unsafe_allow_html=True)
-            if st.button("🔄 Restart Simulation"):
-                for key in [
-                    "annotated_images", "counts", "sorted_directions",
-                    "current_index", "phase", "finished"
-                ]:
-                    if key in st.session_state:
-                        del st.session_state[key]
-                st.rerun()
-                
-            # --- Dashboard ---
+                    
             st.markdown("## 📊 Traffic Analysis Dashboard")
 
             # 1. Total vehicles
-            total_vehicles = sum(st.session_state.counts.values())
-            st.markdown(f"<h3 style='font-size: 24px; font-weight: bold;'>🚗 Total Vehicles Detected: {total_vehicles}</h3>", unsafe_allow_html=True)
-
+            total_vehicles = sum(st.session_state.counts.values())                
+            st.markdown(
+                            f"""
+                            <br>
+                            <h3 style='font-size: 24px; font-weight: bold;'>🚗 Total Vehicles Detected: <span style='color: #4da6ff;'>{total_vehicles}</span></h3>""",unsafe_allow_html=True)
+                
             # 2. Bar chart for vehicle count
-            st.markdown("<h4 style='font-size: 20px; font-weight: bold;'>🚦 Vehicle Count per Direction</h4>", unsafe_allow_html=True)
-            st.bar_chart(st.session_state.counts)
+            st.markdown("<br><h3 style='font-size: 24px; font-weight: bold;'>🚦 Vehicle Count per Direction</h3>", unsafe_allow_html=True)
+            # Create the bar chart
+            vehicle_counts = st.session_state.counts                
+            fig, ax = plt.subplots(figsize=(4,4))
+
+            # Plot the bar char
+            ax.bar(vehicle_counts.keys(), vehicle_counts.values(), color='teal')
+
+            # Add numbers on top of the bars
+            for i, count in enumerate(vehicle_counts.values()):                  
+                ax.text(i, count + 0.5, str(count), ha='center', va='bottom', fontweight='bold', fontsize=8)
+
+            # Customize the chart
+            ax.set_ylabel('Vehicle Count')
+            ax.set_xlabel('Directions')
+            ax.set_title('Vehicle Count per Direction')           
+            # Display the chart in Streamlit
+            # Show in small column
+            col1, col2, _ = st.columns([.3, .4, .3])
+            with col2:
+                st.pyplot(fig)
 
             # 3. Most congested direction
             busiest = max(st.session_state.counts.items(), key=lambda x: x[1])
-            st.markdown(f"<h4 style='font-size: 20px; font-weight: bold;'>**🥇 Busiest Direction:** {busiest[0]} with {busiest[1]} vehicles</h4>", unsafe_allow_html=True)
+            st.markdown(
+                            f"""
+                            <br>
+                            <h3 style='font-size: 24px; font-weight: bold;'>
+                                🥇 Busiest Direction: <span style='color: red;'>{busiest[0]}</span> with {busiest[1]} vehicles
+                            </h3>
+                            """,
+                            unsafe_allow_html=True
+                        )
 
-            # 4. Green time per direction
+                # 4. Green time per direction
             base_time = 5
             time_per_vehicle = 1
-            max_time = 30
-
-            green_times = {
-                direction: min(base_time + count * time_per_vehicle, max_time)
-                for direction, count in st.session_state.counts.items()
-            }
-            st.markdown("<h4 style='font-size: 20px; font-weight: bold;'>⏱️ Assigned Green Time</h4>" , unsafe_allow_html=True)
-            st.markdown("""
-                        <style>
-                        .dataframe th, .dataframe td {
-                        font-size: 14px;
-                        padding: 8px;
-                        text-align: center;
-                        }
-                        .dataframe {
-                        width: 60%;
-                        margin: 0 auto;
-                        }
-                        </style>
-                        """, unsafe_allow_html=True)
-            st.table(green_times)
+            max_time = 25                
+            # Use the same formula applied earlier
+            green_times = [
+                    min(base_time + int(st.session_state.counts[d] / 2) * time_per_vehicle, max_time)
+                    for d, _ in st.session_state.sorted_directions
+                ]
+            green_df = pd.DataFrame({
+                            "Direction": [d for d, _ in sorted_directions],
+                            "Assigned Green Time (sec)": green_times
+                })
+            st.markdown("<br><h3 style='font-size: 24px; font-weight: bold;'>⏱️ Assigned Green Time</h3>" , unsafe_allow_html=True)
+            col1, col2, _ = st.columns([.3, .4, .3])
+            with col2:
+                    st.dataframe(green_df.style.set_properties(**{
+                            'background-color': '#e8f5e9',
+                            'color': 'green',
+                            'font-size': '12px',
+                            'font-weight': 'bold',
+                            'text-align' : 'left'
+                            }), use_container_width=True)
 
             # 5. Pie chart (optional)
             import pandas as pd
             df = pd.DataFrame.from_dict(st.session_state.counts, orient='index', columns=['Vehicles'])
             st.markdown("<h4 style='font-size: 20px; font-weight: bold;'>📈 Traffic Share by Direction</h4>", unsafe_allow_html=True)
-            st.pyplot(df.plot.pie(y='Vehicles', autopct='%1.1f%%', legend=False, figsize=(3,3)).figure)
+            fig2, ax2 = plt.subplots(figsize=(4, 4))  # Smaller size
+            ax2.pie(counts.values(), labels=counts.keys(), autopct='%1.1f%%', textprops={'fontsize': 8})               
+            ax2.set_title("Traffic Distribution", fontsize=10)
+            col1, col2, _ = st.columns([.2, .4, .2])
+            with col2:
+                    st.pyplot(fig2)
+                    
+            # 6. Waiting Time Calculation ---
+            st.markdown("<br>", unsafe_allow_html=True)
+            st.markdown("### ⏱️ Waiting Time Analysis")
+            # Use the already computed busiest direction
+            busiest_direction = busiest[0]
+
+            # Define green time per vehicle
+            time_per_vehicle = 1
+            # Initialize dictionary to store waiting times
+            waiting_times = {}
+
+            # Calculate waiting times for all directions except the busiest one
+            for direction in directions:
+                    if direction == busiest_direction:
+                        waiting_times[direction] = 0  # No wait for the first green
+                    else:
+                        # Waiting time is the sum of green times of all directions before this one (excluding busiest and current)
+                        preceding_directions = [d for d, _ in st.session_state.sorted_directions if d != direction and d != busiest_direction]
+                        wait = sum(min(5 + st.session_state.counts[d] * time_per_vehicle, 30) + 3  # green + yellow
+                                for d in preceding_directions)
+                        waiting_times[direction] = wait
+
+                # Display waiting times for each direction
+            waiting_time_data = pd.DataFrame({
+                    "Direction": directions,
+                    "Waiting Time (sec)": [waiting_times[direction] for direction in directions]
+            })
+
+            # Display waiting times in a table
+            # Show in small column
+            col1, col2, _ = st.columns([.3, .4, .3])
+            with col2:
+                    st.dataframe(waiting_time_data.style.set_properties(**{
+                    'background-color': '#f3f4f6',
+                    'color': 'black',
+                    'font-size': '12px',
+                    'font-weight': 'bold'
+                    }))
+
+            # Optionally, add a message for the waiting times
+            st.markdown("<br>", unsafe_allow_html=True)
+            st.markdown(
+                    """
+                    The waiting times represent the total waiting time for vehicles at each direction based on the vehicle counts in other directions.
+                    The higher the vehicle count in other directions, the higher the waiting time for the current direction.
+                    """
+                )
+            st.session_state.clear()