Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,133 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import cv2
+from ultralytics import YOLO
+import tempfile
+import time
+from crewai import Agent, Task, Crew
+from langchain_groq import ChatGroq
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# --- CONFIG ---
+st.set_page_config(page_title="CityFlow AI", page_icon="🚦", layout="wide")
+
+# --- SIDEBAR ---
+with st.sidebar:
+    st.header("🚦 CityFlow Control")
+    groq_api_key = st.text_input("Groq API Key", type="password")
+    uploaded_file = st.file_uploader("Upload CCTV Footage", type=['mp4', 'mov'])
+    confidence = st.slider("Detection Confidence", 0.0, 1.0, 0.3)
+
+# --- MAIN APP ---
+st.title("🚦 CityFlow: Autonomous Traffic Management")
+st.markdown("Real-time computer vision + Multi-Agent reasoning to optimize traffic flow.")
+
+# Initialize Model (YOLOv8n is small and fast)
+@st.cache_resource
+def load_model():
+    return YOLO('yolov8n.pt')
+
+model = load_model()
+
+# --- THE AGENTS ---
+if groq_api_key:
+    llm = ChatGroq(temperature=0, model_name="llama-3.3-70b-versatile", groq_api_key=groq_api_key)
+
+    # Agent 1: The Traffic Analyst (Reads the data)
+    analyst = Agent(
+        role='Traffic Data Analyst',
+        goal='Analyze vehicle counts and congestion levels to determine traffic severity.',
+        backstory="You are an expert in urban flow. You look at raw numbers (car counts) and decide if it's 'Light', 'Heavy', or 'Gridlock'.",
+        llm=llm,
+        verbose=False
+    )
+
+    # Agent 2: The Signal Controller (Makes the decision)
+    controller = Agent(
+        role='Signal Control Officer',
+        goal='Decide the optimal traffic light duration based on severity.',
+        backstory="You control the city lights. If an ambulance is seen, you MUST open the lane. If traffic is heavy, extend Green light duration.",
+        llm=llm,
+        verbose=False
+    )
+
+# --- VIDEO PROCESSING LOOP ---
+if uploaded_file and groq_api_key:
+    # Save temp file
+    tfile = tempfile.NamedTemporaryFile(delete=False)
+    tfile.write(uploaded_file.read())
+    cap = cv2.VideoCapture(tfile.name)
+
+    col1, col2 = st.columns([2, 1])
+    
+    with col1:
+        st_frame = st.empty()
+    
+    with col2:
+        st_metrics = st.empty()
+        st_decision = st.empty()
+
+    frame_count = 0
+    
+    # Process Video
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        frame_count += 1
+        
+        # Only run AI every 10 frames to keep it fast
+        if frame_count % 10 == 0:
+            
+            # 1. PERCEPTION (YOLO)
+            results = model.track(frame, persist=True, conf=confidence)
+            
+            # Count Vehicles
+            car_count = 0
+            emergency_count = 0
+            
+            # Classes: 2=Car, 3=Motorcycle, 5=Bus, 7=Truck
+            # We assume class 0 is person, so we filter for vehicles.
+            # Note: Standard COCO dataset doesn't distinguish "ambulance", 
+            # so we'll simulate emergency logic if we see a Truck (class 7) for this demo.
+            
+            for result in results:
+                boxes = result.boxes
+                for box in boxes:
+                    cls = int(box.cls[0])
+                    if cls in [2, 3, 5, 7]: 
+                        car_count += 1
+                    if cls == 7: # Simulating "Truck/Emergency" for demo logic
+                        emergency_count += 1
+
+            # Draw Boxes on Frame
+            res_plotted = results[0].plot()
+            
+            # 2. REASONING (CrewAI)
+            # We pass the real-time data to the agents
+            
+            task_analyze = Task(
+                description=f"Current Status: {car_count} vehicles detected. {emergency_count} heavy/emergency vehicles. Analyze congestion level.",
+                agent=analyst,
+                expected_output="A short status: 'Light', 'Moderate', 'Critical'."
+            )
+            
+            task_control = Task(
+                description="Based on the analysis, decide the Green Light duration (in seconds). If Status is Critical or Emergency detected, set to MAX (60s).",
+                agent=controller,
+                expected_output="JSON: {'status': '...', 'green_light_duration': 45, 'reason': '...'}"
+            )
+
+            # Creating a mini-crew for this single frame decision
+            # (In production, you'd run this async, not blocking the video)
+            crew = Crew(agents=[analyst, controller], tasks=[task_analyze, task_control])
+            decision = crew.kickoff()
+
+            # 3. VISUALIZATION
+            st_frame.image(res_plotted, channels="BGR")
+            
+            st_metrics.metric(label="Vehicles Detected", value=car_count)
+            st_metrics.metric(label="Emergency/Heavy", value=emergency_count)
+            
+            st_decision.info(f"🤖 AI Decision: \n\n{decision}")
+
+    cap.release()