Spaces:
Sleeping
Sleeping
File size: 4,829 Bytes
b321750 1b3baae b321750 1b3baae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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
# --- 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() |