Spaces:
Running
Running
| from ultralytics import YOLO | |
| import base64 | |
| import cv2 | |
| import io | |
| import numpy as np | |
| from ultralytics.utils.plotting import Annotator | |
| import streamlit as st | |
| from streamlit_image_coordinates import streamlit_image_coordinates | |
| import pandas as pd | |
| import ollama | |
| import bs4 | |
| import tempfile | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_community.document_loaders import WebBaseLoader | |
| from langchain_community.document_loaders import CSVLoader | |
| from langchain_community.vectorstores import Chroma | |
| from langchain_community.embeddings import OllamaEmbeddings | |
| from langchain_core.output_parsers import StrOutputParser | |
| from langchain_core.runnables import RunnablePassthrough | |
| def set_background(image_file1,image_file2): | |
| with open(image_file1, "rb") as f: | |
| img_data1 = f.read() | |
| b64_encoded1 = base64.b64encode(img_data1).decode() | |
| with open(image_file2, "rb") as f: | |
| img_data2 = f.read() | |
| b64_encoded2 = base64.b64encode(img_data2).decode() | |
| style = f""" | |
| <style> | |
| .stApp{{ | |
| background-image: url(data:image/png;base64,{b64_encoded1}); | |
| background-size: cover; | |
| }} | |
| .st-emotion-cache-6qob1r{{ | |
| background-image: url(data:image/png;base64,{b64_encoded2}); | |
| background-size: cover; | |
| border: 5px solid rgb(14, 17, 23); | |
| }} | |
| </style> | |
| """ | |
| st.markdown(style, unsafe_allow_html=True) | |
| set_background('pngtree-city-map-navigation-interface-picture-image_1833642.png','2024-05-18_14-57-09_5235.png') | |
| st.title("Traffic Flow and Optimization Toolkit") | |
| sb = st.sidebar # defining the sidebar | |
| sb.markdown("🛰️ **Navigation**") | |
| page_names = ["PS1", "PS2", "PS3","Chat with Results"] | |
| page = sb.radio("", page_names, index=0) | |
| st.session_state['n'] = sb.slider("Number of ROIs",1,5) | |
| if page == 'PS1': | |
| uploaded_file = st.file_uploader("Choose a video...", type=["mp4", "mpeg"]) | |
| model = YOLO('yolov8n.pt') | |
| if uploaded_file is not None: | |
| with tempfile.NamedTemporaryFile(delete=False) as temp: | |
| temp.write(uploaded_file.read()) | |
| if 'roi_list1' not in st.session_state: | |
| st.session_state['roi_list1'] = [] | |
| if "all_rois1" not in st.session_state: | |
| st.session_state['all_rois1'] = [] | |
| classes = model.names | |
| done_1 = st.button('Selection Done') | |
| while len(st.session_state["all_rois1"]) < st.session_state['n']: | |
| cap = cv2.VideoCapture(temp.name) | |
| while not done_1: | |
| ret,frame=cap.read() | |
| cv2.putText(frame,'SELECT ROI',(100,100),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),4) | |
| if not ret: | |
| st.write('ROI selection unsuccessfull') | |
| break | |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| value = streamlit_image_coordinates(frame,key='numpy',width=750) | |
| st.session_state["roi_list1"].append([int(value['x']*2.55),int(value['y']*2.55)]) | |
| st.write(st.session_state["roi_list1"]) | |
| if cv2.waitKey(0)&0xFF==27: | |
| break | |
| cap.release() | |
| st.session_state["all_rois1"].append(st.session_state["roi_list1"]) | |
| st.session_state["roi_list1"] = [] | |
| done_1 = False | |
| st.write('ROI indices: ',st.session_state["all_rois1"][0]) | |
| cap = cv2.VideoCapture(temp.name) | |
| st.write("Detection started") | |
| st.session_state['fps'] = cap.get(cv2.CAP_PROP_FPS) | |
| st.write(f"FPS OF VIDEO: {st.session_state['fps']}") | |
| avg_list = [] | |
| count = 0 | |
| frame_placeholder = st.empty() | |
| st.session_state["data1"] = {} | |
| for i in range(len(st.session_state["all_rois1"])): | |
| st.session_state["data1"][f"ROI{i}"] = [] | |
| while cap.isOpened(): | |
| ret,frame=cap.read() | |
| if not ret: | |
| break | |
| count += 1 | |
| if count % 3 != 0: | |
| continue | |
| k = 0 | |
| for roi_list_here1 in st.session_state["all_rois1"]: | |
| max = [0,0] | |
| min = [10000,10000] | |
| roi_list_here = roi_list_here1[1:] | |
| for i in range(len(roi_list_here)): | |
| if roi_list_here[i][0] > max[0]: | |
| max[0] = roi_list_here[i][0] | |
| if roi_list_here[i][1] > max[1]: | |
| max[1] = roi_list_here[i][1] | |
| if roi_list_here[i][0] < min[0]: | |
| min[0] = roi_list_here[i][0] | |
| if roi_list_here[i][1] < min[1]: | |
| min[1] = roi_list_here[i][1] | |
| frame_cropped = frame[min[1]:max[1],min[0]:max[0]] | |
| roi_corners = np.array([roi_list_here],dtype=np.int32) | |
| mask = np.zeros(frame.shape,dtype=np.uint8) | |
| mask.fill(255) | |
| channel_count = frame.shape[2] | |
| ignore_mask_color = (255,)*channel_count | |
| cv2.fillPoly(mask,roi_corners,0) | |
| mask_cropped = mask[min[1]:max[1],min[0]:max[0]] | |
| roi = cv2.bitwise_or(frame_cropped,mask_cropped) | |
| #roi = frame[roi_list_here[0][1]:roi_list_here[1][1],roi_list_here[0][0]:roi_list_here[1][0]] | |
| number = [] | |
| results = model.predict(roi) | |
| for r in results: | |
| boxes = r.boxes | |
| counter = 0 | |
| for box in boxes: | |
| counter += 1 | |
| name = classes[box.cls.numpy()[0]] | |
| conf = str(round(box.conf.numpy()[0],2)) | |
| text = name+""+conf | |
| bbox = box.xyxy[0].numpy() | |
| cv2.rectangle(frame,(int(bbox[0])+min[0],int(bbox[1])+min[1]),(int(bbox[2])+min[0],int(bbox[3])+min[1]),(0,255,0),2) | |
| cv2.putText(frame,text,(int(bbox[0])+min[0],int(bbox[1])+min[1]-5),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),2) | |
| number.append(counter) | |
| avg = sum(number)/len(number) | |
| stats = str(round(avg,2)) | |
| if count%10 == 0: | |
| st.session_state["data1"][f"ROI{k}"].append(avg) | |
| k+=1 | |
| cv2.putText(frame,stats,(min[0],min[1]),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,0),4) | |
| cv2.polylines(frame,roi_corners,True,(255,0,0),2) | |
| cv2.putText(frame,'The average number of vehicles in the Regions of Interest',(100,100),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),4) | |
| frame_placeholder.image(frame,channels='BGR') | |
| cap.release() | |
| st.write("The resultant data is:") | |
| st.write(st.session_state.data1) | |
| else: | |
| st.error('PLEASE UPLOAD', icon="🚨") | |
| elif page == "PS3": | |
| uploaded_file1 = st.file_uploader("Choose a video...", type=["mp4", "mpeg"]) | |
| model1 = YOLO("yolov8n.pt") | |
| model2 = YOLO("best.pt") | |
| if uploaded_file1 is not None: | |
| with tempfile.NamedTemporaryFile(delete=False) as temp: | |
| temp.write(uploaded_file1.read()) | |
| if 'roi_list2' not in st.session_state: | |
| st.session_state['roi_list2'] = [] | |
| if "all_rois2" not in st.session_state: | |
| st.session_state['all_rois2'] = [] | |
| classes = model1.names | |
| done_2 = st.button('Selection Done') | |
| while len(st.session_state["all_rois2"]) < st.session_state['n']: | |
| cap = cv2.VideoCapture(temp.name) | |
| while not done_2: | |
| ret,frame=cap.read() | |
| cv2.putText(frame,'SELECT ROI',(100,100),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),4) | |
| if not ret: | |
| st.write('ROI selection has concluded') | |
| break | |
| frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| value = streamlit_image_coordinates(frame,key='numpy',width=750) | |
| st.session_state["roi_list2"].append([int(value['x']*2.5),int(value['y']*2.5)]) | |
| st.write(st.session_state["roi_list2"]) | |
| if cv2.waitKey(0)&0xFF==27: | |
| break | |
| cap.release() | |
| st.session_state["all_rois2"].append(st.session_state["roi_list2"]) | |
| st.session_state["roi_list2"] = [] | |
| done_2 = False | |
| st.write('ROI indices: ',st.session_state["all_rois2"][0]) | |
| cap = cv2.VideoCapture(temp.name) | |
| st.write("Detection started") | |
| avg_list = [] | |
| count = 0 | |
| frame_placeholder = st.empty() | |
| st.session_state.data = {} | |
| for i in range(len(st.session_state["all_rois2"])): | |
| st.session_state["data"][f"ROI{i}"] = [] | |
| for i in range(len(st.session_state['all_rois2'])): | |
| st.session_state.data[f"ROI{i}"] = [] | |
| while cap.isOpened(): | |
| ret,frame=cap.read() | |
| if not ret: | |
| break | |
| count += 1 | |
| if count % 3 != 0: | |
| continue | |
| # rois = [] | |
| k = 0 | |
| for roi_list_here1 in st.session_state["all_rois2"]: | |
| max = [0,0] | |
| min = [10000,10000] | |
| roi_list_here = roi_list_here1[1:] | |
| for i in range(len(roi_list_here)-1): | |
| if roi_list_here[i][0] > max[0]: | |
| max[0] = roi_list_here[i][0] | |
| if roi_list_here[i][1] > max[1]: | |
| max[1] = roi_list_here[i][1] | |
| if roi_list_here[i][0] < min[0]: | |
| min[0] = roi_list_here[i][0] | |
| if roi_list_here[i][1] < min[1]: | |
| min[1] = roi_list_here[i][1] | |
| frame_cropped = frame[min[1]:max[1],min[0]:max[0]] | |
| roi_corners = np.array([roi_list_here],dtype=np.int32) | |
| mask = np.zeros(frame.shape,dtype=np.uint8) | |
| mask.fill(255) | |
| channel_count = frame.shape[2] | |
| ignore_mask_color = (255,)*channel_count | |
| cv2.fillPoly(mask,roi_corners,0) | |
| mask_cropped = mask[min[1]:max[1],min[0]:max[0]] | |
| roi = cv2.bitwise_or(frame_cropped,mask_cropped) | |
| #roi = frame[roi_list_here[0][1]:roi_list_here[1][1],roi_list_here[0][0]:roi_list_here[1][0]] | |
| number = [] | |
| results = model1.predict(roi) | |
| results_pothole = model2.predict(source=frame) | |
| for r in results: | |
| boxes = r.boxes | |
| counter = 0 | |
| for box in boxes: | |
| counter += 1 | |
| name = classes[box.cls.numpy()[0]] | |
| conf = str(round(box.conf.numpy()[0],2)) | |
| text = name+conf | |
| bbox = box.xyxy[0].numpy() | |
| cv2.rectangle(frame,(int(bbox[0])+min[0],int(bbox[1])+min[1]),(int(bbox[2])+min[0],int(bbox[3])+min[1]),(0,255,0),2) | |
| cv2.putText(frame,text,(int(bbox[0])+min[0],int(bbox[1])+min[1]-5),cv2.FONT_HERSHEY_SIMPLEX, 0.4,(0,0,255),2) | |
| number.append(counter) | |
| for r in results_pothole: | |
| masks = r.masks | |
| boxes = r.boxes.cpu().numpy() | |
| xyxys = boxes.xyxy | |
| confs = boxes.conf | |
| if masks is not None: | |
| shapes = np.ones_like(frame) | |
| for mask,conf,xyxy in zip(masks,confs,xyxys): | |
| polygon = mask.xy[0] | |
| if conf >= 0.49 and len(polygon)>=3: | |
| cv2.fillPoly(shapes,pts=np.int32([polygon]),color=(0,0,255,0.5)) | |
| frame = cv2.addWeighted(frame,0.7,shapes,0.3,gamma=0) | |
| cv2.rectangle(frame,(int(xyxy[0]),int(xyxy[1])),(int(xyxy[2]),int(xyxy[3])),(0,0,255),2) | |
| cv2.putText(frame,'Pothole '+str(conf),(int(xyxy[0]),int(xyxy[1])-5),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),2) | |
| avg = sum(number)/len(number) | |
| stats = str(round(avg,2)) | |
| if count % 10 == 0: | |
| st.session_state.data[f"ROI{k}"].append(avg) | |
| k+=1 | |
| cv2.putText(frame,stats,(min[0],min[1]),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,0),4) | |
| cv2.polylines(frame,roi_corners,True,(255,0,0),2) | |
| if counter >= 5: | |
| cv2.putText(frame,'!!CONGESTION MORE THAN '+str(counter)+' Objects',(min[0]+20,min[1]+20),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,0),4) | |
| cv2.polylines(frame,roi_corners,True,(255,0,0),2) | |
| cv2.putText(frame,'Objects in the Regions of Interest',(100,100),cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),4) | |
| frame_placeholder.image(frame,channels='BGR') | |
| cap.release() | |
| st.write("The result is:") | |
| st.write(st.session.data) | |
| else: | |
| st.error('PLEASE UPLOAD', icon="🚨") | |
| elif page == "PS2": | |
| st.header("CLICK ON RUN SCRIPT TO START A TRAFFIC SIMULATION") | |
| script = st.button("RUN SCRIPT") | |
| st.error("This simulation, Unfortunately was not optimized for deployed but still works for locally run version", icon="🚨") | |
| st.write("check out this [link](https://github.com/pranavsrinivasa/Traffic-Optimization-Projects)") | |
| elif page == "Chat with Results": | |
| st.title('Chat with the Results') | |
| st.write("Please upload the relevant CSV data to get started") | |
| st.error("This model, Unfortunately was not optimized for deployed but still works for locally run version", icon="🚨") | |
| st.write("check out this [link](https://github.com/pranavsrinivasa/Traffic-Optimization-Projects)") | |