Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cv2 | |
| import tempfile | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from huggingface_hub import hf_hub_download | |
| # Force the page to wide layout for better visibility | |
| st.set_page_config(page_title="VirtuGuard AI Security Suite 2026", layout="wide") | |
| # 2. Add your logo | |
| # This renders it in the upper-left corner and sidebar | |
| st.logo("logo.png", size="medium") | |
| # Optional: Add a larger version or specific text to the sidebar | |
| with st.sidebar: | |
| st.image("logo.png", width=150) | |
| st.markdown("VirtuGuard AI Security Suite 2026") | |
| # --- Load Models --- | |
| def load_models(): | |
| # 1. Advanced YOLO26 for Person (0) and Vehicles (2,3,5,7) | |
| main = YOLO('yolo26n.pt') | |
| # 2. Corrected Firearm/Threat model path | |
| # 'weights/best.pt' is the correct location for this repo | |
| threat_path = hf_hub_download( | |
| repo_id="Subh775/Threat-Detection-YOLOv8n", | |
| filename="weights/best.pt" | |
| ) | |
| threat = YOLO(threat_path) | |
| return main, threat | |
| main_model, threat_model = load_models() | |
| st.title("🛡️ VirtuGuard Advanced AI Security Suite (YOLO26)") | |
| st.info("Detecting People, Vehicles, and Multi-Threats (Guns & Knives)") | |
| uploaded_file = st.file_uploader("Upload Image or Video", type=["mp4", "mov", "jpg", "png"]) | |
| if uploaded_file: | |
| is_video = uploaded_file.type.startswith('video') | |
| if is_video: | |
| # Create a temp file to store the uploaded video | |
| tfile = tempfile.NamedTemporaryFile(delete=False) | |
| tfile.write(uploaded_file.read()) | |
| cap = cv2.VideoCapture(tfile.name) | |
| st_frame = st.empty() | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: break | |
| # Inference | |
| # Main model: Person=0, Car=2, Motorcycle=3, Bus=5, Truck=7 | |
| res_main = main_model(frame, classes=[0, 2, 3, 5, 7], verbose=False)[0] | |
| # Threat model: 1=Gun, 2=Explosive, 3=Grenade, 4=Knife | |
| res_threat = threat_model(frame, conf=0.5, verbose=False)[0] | |
| # Combine visualizations | |
| annotated = res_main.plot() | |
| annotated = res_threat.plot(img=annotated) | |
| # Display | |
| st_frame.image(annotated, channels="BGR") | |
| cap.release() | |
| else: | |
| # Static Image Processing | |
| img = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1) | |
| r1 = main_model(img, classes=[0, 2, 3, 5, 7])[0] | |
| r2 = threat_model(img, conf=0.5)[0] | |
| # Merge results onto one image | |
| final_img = r2.plot(img=r1.plot()) | |
| st.image(final_img, channels="BGR", use_container_width=True) | |
| # Display Summary | |
| st.subheader("Detections Summary") | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.write("**Base Detections:**") | |
| for box in r1.boxes: | |
| label = main_model.names[int(box.cls[0])] | |
| st.write(f"✅ {label.capitalize()}: {box.conf[0]:.2%}") | |
| with col2: | |
| st.write("**Security Threats:**") | |
| for box in r2.boxes: | |
| # Map threat model class IDs to names | |
| threat_names = {1: "Gun", 2: "Explosive", 3: "Grenade", 4: "Knife"} | |
| label = threat_names.get(int(box.cls[0]), "Weapon") | |
| st.write(f"⚠️ {label.upper()}: {box.conf[0]:.2%}") |