Update app.py

app.py

@@ -1,91 +1,158 @@
-import sys
-from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QTextEdit, QPushButton, QLineEdit, QCheckBox, QHBoxLayout
-from PyQt5.QtGui import QFont, QColor, QPalette
-from PyQt5.QtCore import Qt
-from huggingface_hub import InferenceClient
-
-class VarunGPTApp(QWidget):
-    def __init__(self):
-        super().__init__()
-        self.initUI()
-        self.chat_history = ""
-    
-    def initUI(self):
-        self.setWindowTitle("VarunGPT - AI Chatbot")
-        self.setGeometry(200, 200, 700, 550)
-        
-        palette = self.palette()
-        palette.setColor(QPalette.Window, QColor(20, 20, 20))
-        self.setPalette(palette)
-        
-        layout = QVBoxLayout()
-        
-        self.title = QLabel("VarunGPT", self)
-        self.title.setFont(QFont("Arial", 22, QFont.Bold))
-        self.title.setStyleSheet("color: #ff9800; text-align: center;")
-        self.title.setAlignment(Qt.AlignCenter)
-        layout.addWidget(self.title)
-        
-        self.api_key_input = QLineEdit(self)
-        self.api_key_input.setPlaceholderText("Enter API Key")
-        self.api_key_input.setEchoMode(QLineEdit.Password)
-        self.api_key_input.setStyleSheet("background-color: #222; color: white; padding: 8px; border-radius: 5px;")
-        layout.addWidget(self.api_key_input)
-        
-        self.theme_checkbox = QCheckBox("Enable Advanced Reasoning")
-        self.theme_checkbox.setStyleSheet("color: white;")
-        layout.addWidget(self.theme_checkbox)
-        
-        self.user_input = QLineEdit(self)
-        self.user_input.setPlaceholderText("Ask anything...")
-        self.user_input.setStyleSheet("background-color: #222; color: white; padding: 8px; border-radius: 5px;")
-        layout.addWidget(self.user_input)
-        
-        self.send_button = QPushButton("Send", self)
-        self.send_button.setStyleSheet("background-color: #ff9800; color: white; padding: 10px; border-radius: 5px;")
-        self.send_button.clicked.connect(self.handleUserInput)
-        layout.addWidget(self.send_button)
-        
-        self.chat_display = QTextEdit(self)
-        self.chat_display.setReadOnly(True)
-        self.chat_display.setStyleSheet("background-color: #111; color: white; padding: 10px; border-radius: 5px;")
-        layout.addWidget(self.chat_display)
-        
-        self.setLayout(layout)
-    
-    def handleUserInput(self):
-        api_key = self.api_key_input.text().strip()
-        if not api_key:
-            self.chat_display.setText("Error: API Key is required.")
-            return
-        
-        client = InferenceClient(api_key=api_key)
-        
-        user_text = self.user_input.text()
-        if user_text:
-            self.chat_history += f"You: {user_text}\n"
-            self.chat_display.setText(self.chat_history)
-            
-            ai_response = self.query_huggingface(client, user_text)
-            self.chat_history += f"VarunGPT: {ai_response}\n"
-            self.chat_display.setText(self.chat_history)
-            
-            self.user_input.clear()
-    
-    def query_huggingface(self, client, prompt):
-        messages = [{"role": "user", "content": prompt}]
-        try:
-            completion = client.chat.completions.create(
-                model="mistralai/Mistral-7B-Instruct-v0.3", 
-                messages=messages, 
-                max_tokens=500
-            )
-            return completion["choices"][0]["message"]["content"]
-        except Exception as e:
-            return f"Error: {str(e)}"
-
-if __name__ == "__main__":
-    app = QApplication(sys.argv)
-    window = VarunGPTApp()
-    window.show()
-    sys.exit(app.exec_())
+import streamlit as st
+import cv2
+from PIL import Image
+import numpy as np
+from ultralytics import YOLO
+import tempfile
+import os
+import supervision as sv
+from streamlit_webrtc import webrtc_streamer, WebRtcMode
+import functools
+
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="Traffic Lane and Object Detection",
+    page_icon=":camera_video:",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+# --- Sidebar ---
+st.sidebar.header("Traffic Lane and Object Detection Options")
+source_type = st.sidebar.radio("Select Input Source:", ("Image", "Video", "Live Camera Feed")) # Added Live Camera Feed
+confidence_threshold = st.sidebar.slider(
+    "Confidence Threshold", min_value=0.0, max_value=1.0, value=0.25, step=0.05
+)
+iou_threshold = st.sidebar.slider(
+    "IoU Threshold", min_value=0.0, max_value=1.0, value=0.45, step=0.05, help="Intersection over Union threshold for NMS"
+)
+
+
+# --- Load YOLO Model ---
+@st.cache_resource
+def load_model():
+    model = YOLO("yolov8n.pt")  # Load a pretrained YOLOv8n model
+    return model
+
+model = load_model()
+
+# --- Functions ---
+def detect_lanes_and_objects_image(image, model, confidence_threshold, iou_threshold):
+    """Runs YOLO on a single image with Supervision."""
+    img_np = np.array(image)
+    results = model(img_np, conf=confidence_threshold, iou=iou_threshold)
+
+    detections = sv.Detections.from_ultralytics(results[0])
+    annotator = sv.BoxAnnotator(thickness=2, text_thickness=1, text_scale=0.5)
+    annotated_frame = annotator.annotate(scene=img_np, detections=detections)
+
+    return annotated_frame
+
+
+def detect_lanes_and_objects_video(video_path, model, confidence_threshold, iou_threshold):
+    """Runs YOLO on a video file with Supervision."""
+    video = cv2.VideoCapture(video_path)
+    frame_width = int(video.get(3))
+    frame_height = int(video.get(4))
+    fps = int(video.get(cv2.CAP_PROP_FPS))
+    codec = cv2.VideoWriter_fourcc(*"mp4v")
+    output_path = "output.mp4"
+    out = cv2.VideoWriter(output_path, codec, fps, (frame_width, frame_height))
+    stframe = st.empty()
+
+    while video.isOpened():
+        ret, frame = video.read()
+        if not ret:
+            break
+
+        results = model(frame, conf=confidence_threshold, iou=iou_threshold)
+
+        detections = sv.Detections.from_ultralytics(results[0])
+        annotator = sv.BoxAnnotator(thickness=2, text_thickness=1, text_scale=0.5)
+        annotated_frame = annotator.annotate(scene=frame, detections=detections)
+
+        out.write(annotated_frame)
+        stframe.image(annotated_frame, channels="BGR", use_column_width=True)
+
+    video.release()
+    out.release()
+    cv2.destroyAllWindows()
+
+    return output_path
+
+
+def process_frame(frame, model, confidence_threshold, iou_threshold):
+    """Process each frame from the webcam with Supervision."""
+    img = frame.to_ndarray(format="bgr24")
+    results = model(img, conf=confidence_threshold, iou=iou_threshold)
+
+    detections = sv.Detections.from_ultralytics(results[0])
+    annotator = sv.BoxAnnotator(thickness=2, text_thickness=1, text_scale=0.5)
+    annotated_frame = annotator.annotate(scene=img, detections=detections)
+
+    return annotated_frame
+
+
+# --- Main Application ---
+st.title("Traffic Lane and Object Detection")
+
+if source_type == "Image":
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+
+        if st.button("Run Detection"):
+            with st.spinner("Running YOLOv8..."):
+                detected_image = detect_lanes_and_objects_image(
+                    image, model, confidence_threshold, iou_threshold
+                )
+                st.image(detected_image, caption="Detected Image", use_column_width=True)
+
+elif source_type == "Video":
+    uploaded_video = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"])
+
+    if uploaded_video is not None:
+        tfile = tempfile.NamedTemporaryFile(delete=False)
+        tfile.write(uploaded_video.read())
+        video_path = tfile.name
+
+        if st.button("Run Detection"):
+            with st.spinner("Running YOLOv8 on video..."):
+                output_video_path = detect_lanes_and_objects_video(
+                    video_path, model, confidence_threshold, iou_threshold
+                )
+                st.video(output_video_path)
+                tfile.close()
+                os.unlink(video_path)
+                os.remove(output_video_path)
+
+elif source_type == "Live Camera Feed":
+    st.subheader("Live Camera Feed")
+
+    if model is None:
+        st.error("YOLO model is not loaded. Cannot run live feed.")
+        st.stop()
+
+    custom_process_frame = functools.partial(
+        process_frame,
+        model=model,
+        confidence_threshold=confidence_threshold,
+        iou_threshold=iou_threshold,
+    )
+
+    webrtc_streamer(
+        key="live-feed",
+        video_frame_callback=custom_process_frame,
+        mode=WebRtcMode.SENDRECV,
+        media_stream_constraints={"video": True, "audio": False},
+    )
+
+
+st.markdown("---")
+st.markdown(
+    """
+    **Note:** This example uses YOLOv8 for object detection. Lane detection is a more complex task and requires additional image processing techniques. This is a simplified demo and will likely not perform well on complex or noisy video.
+    """
+)