Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,12 +2,62 @@ import altair as alt
 import numpy as np
 import pandas as pd
 import streamlit as st
-
-import streamlit as st
+import cv2
+from PIL import Image
+from ultralytics import YOLO
 from traffic_logic import get_next_green
 
 st.set_page_config(page_title="Smart Traffic Light System", layout="wide")
-st.title("🚦 Smart Traffic Light Simulation")
+# st.title("🚦 Smart Traffic Light Simulation")
+
+# Load YOLO model
+model = YOLO("yolov8n.pt")
+
+# UI
+st.title("🚦 Smart Traffic Light System for 4-Way Intersection (YOLO-Based)")
+st.markdown("Upload traffic images from **4 directions** to simulate smart light control.")  
+# Upload images for 4 directions
+cols = st.columns(4)
+with cols[0]: n_img = st.file_uploader("North", type=["jpg", "png", "jpeg"], key="north")
+with cols[1]: s_img = st.file_uploader("South", type=["jpg", "png", "jpeg"], key="south")
+with cols[2]: e_img = st.file_uploader("East", type=["jpg", "png", "jpeg"], key="east")
+with cols[3]: w_img = st.file_uploader("West", type=["jpg", "png", "jpeg"], key="west")
+images = {'North': n_img, 'South': s_img, 'East': e_img, 'West': w_img}
+counts = {}
+results_imgs = {}
+
+if all(images.values()):
+    for direction, uploaded_file in images.items():
+        img = Image.open(uploaded_file).convert("RGB")
+        np_img = np.array(img)
+        
+        result = model.predict(np_img)[0]
+        detections = result.boxes
+
+        # Count cars, trucks, etc.
+        car_classes = [2, 3, 5, 7]
+        car_count = sum(1 for box in detections if int(box.cls) in car_classes)
+        counts[direction] = car_count
+
+        # Draw boxes
+        results_imgs[direction] = result.plot()
+
+    # Determine highest traffic
+    max_dir = max(counts, key=counts.get)
+
+    st.markdown("### 📸 YOLO Detection Results")
+    for direction in ['North', 'South', 'East', 'West']:
+        st.image(results_imgs[direction], caption=f"{direction} - Vehicles: {counts[direction]}")
+
+    # Show traffic light result
+    st.markdown("## 🚦 Traffic Light Control Result:")
+    for direction in ['North', 'South', 'East', 'West']:
+        if direction == max_dir:
+            st.success(f"🟢 {direction} → GREEN (Traffic: {counts[direction]})")
+        else:
+            st.error(f"🔴 {direction} → RED (Traffic: {counts[direction]})")
+else:
+    st.warning("Please upload images for all 4 directions to simulate.")
 
 st.sidebar.header("Vehicle Count per Road")
 north = st.sidebar.slider("North Road", 0, 50, 10)
@@ -18,9 +68,9 @@ west = st.sidebar.slider("West Road", 0, 50, 15)
 vehicle_counts = {'North': north, 'East': east, 'South': south, 'West': west}
 next_green = get_next_green(vehicle_counts)
 
-st.subheader("Traffic Light Status")
+# st.subheader("Traffic Light Status")
 
-cols = st.columns(4)
-for i, (road, count) in enumerate(vehicle_counts.items()):
-    light_color = "🟢 Green" if road == next_green else "🔴 Red"
-    cols[i].metric(label=f"{road} Road", value=f"{count} vehicles", delta=light_color)
+# cols = st.columns(4)
+# for i, (road, count) in enumerate(vehicle_counts.items()):
+#     light_color = "🟢 Green" if road == next_green else "🔴 Red"
+#     cols[i].metric(label=f"{road} Road", value=f"{count} vehicles", delta=light_color)