Update app.py

app.py

@@ -3,130 +3,193 @@ from PIL import Image
 import numpy as np
 import pandas as pd
 import tempfile
+import base64
+import folium
 from reportlab.pdfgen import canvas
 from ultralytics import YOLO
 from transformers import pipeline
-import folium
+import os
 
-# ---------------- Load Models ----------------
-detection_model = YOLO('yolov5s.pt')
+# ------------------------------
+# LOAD MODELS
+# ------------------------------
+detection_model = YOLO("yolov5s.pt")
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-# ---------------- Global History ----------------
-history_df = pd.DataFrame(columns=["Objects", "Severity", "Summary", "Latitude", "Longitude"])
-
-# ---------------- Helper Functions ----------------
-def severity_badge(level):
-    colors = {"Low":"#4CAF50","Medium":"#FF9800","High":"#F44336"}
-    return f"""<div style='padding:8px 15px;border-radius:8px;
-               background:{colors.get(level,"gray")};color:white;
-               width:120px;text-align:center;font-weight:bold;'>{level}</div>"""
-
+# ------------------------------
+# GLOBAL STORAGE
+# ------------------------------
+history = []
+
+# ------------------------------
+# COLOR BADGES
+# ------------------------------
+def severity_badge(sev):
+    colors = {
+        "High": "#ff4b4b",
+        "Medium": "#ffa500",
+        "Low": "#2ecc71"
+    }
+    return f"<span style='padding:6px 12px;border-radius:5px;background:{colors[sev]};color:white;font-weight:bold'>{sev}</span>"
+
+# ------------------------------
+# MAP GENERATOR
+# ------------------------------
 def generate_map(lat, lon):
-    m = folium.Map(location=[lat, lon], zoom_start=14)
+    m = folium.Map(location=[lat, lon], zoom_start=15)
     folium.Marker([lat, lon], popup="Incident Location").add_to(m)
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
-    m.save(tmp.name)
-    return f"<iframe src='{tmp.name}' width='100%' height='350' style='border-radius:8px;border:2px solid #444;'></iframe>"
 
-# ---------------- Main Processing ----------------
-def process_image(img, lat, lon):
-    global history_df
+    f = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
+    m.save(f.name)
+    f.close()
+
+    # Embed map in iframe
+    html = f"""<iframe src="{f.name}" width="100%" height="400" style="border:none;"></iframe>"""
+    return html
+
+# ------------------------------
+# PDF GENERATOR
+# ------------------------------
+def create_pdf(objects, severity, summary, lat, lon):
+    f = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
+    c = canvas.Canvas(f.name)
+
+    c.setFont("Helvetica-Bold", 18)
+    c.drawString(100, 800, "InfraGuard Incident Report")
+
+    c.setFont("Helvetica", 12)
+    c.drawString(100, 770, f"Objects: {objects}")
+    c.drawString(100, 750, f"Severity: {severity}")
+    c.drawString(100, 730, f"AI Summary: {summary}")
+    c.drawString(100, 710, f"Location: {lat}, {lon}")
+
+    c.save()
+    return f.name
+
+# ------------------------------
+# MAIN PROCESS
+# ------------------------------
+def process(img, lat, lon):
     if img is None:
-        return None, "Upload image first!", "", None, "", history_df
+        return None, "Upload Image First", "N/A", None, None, None, history
 
-    # Convert to numpy
-    image_np = np.array(img)
+    img_np = np.array(img)
 
-    # YOLO Prediction
-    results = detection_model.predict(source=image_np, imgsz=640)
-    detected = [int(box.cls) for box in results[0].boxes] if len(results) else []
-    class_names = detection_model.names
-    names = [class_names[c] for c in detected] if detected else ["None"]
+    # YOLO detection
+    results = detection_model.predict(img_np, imgsz=640)
+    objs = []
+    for box in results[0].boxes:
+        objs.append(detection_model.names[int(box.cls)])
 
-    # Severity rules
-    if "fire" in names or (names.count("person") > 10):
+    objects_text = ", ".join(objs) if objs else "None"
+
+    # Severity logic
+    severity = "Low"
+    if "fire" in objs:
         severity = "High"
-    elif "car" in names or "truck" in names:
+    elif "car" in objs or "truck" in objs:
         severity = "Medium"
-    else:
-        severity = "Low"
 
-    # AI Summary
-    text = f"Detected: {', '.join(names)}. Severity is {severity}."
+    # AI summary
+    incident_text = f"Detected objects: {objects_text}. Severity: {severity}"
     try:
-        summary = summarizer(text, max_length=50, min_length=20, do_sample=False)[0]['summary_text']
+        summary = summarizer(incident_text, max_length=40, min_length=10, do_sample=False)[0]['summary_text']
     except:
-        summary = text
-
-    # Update history
-    new_row = pd.DataFrame([[', '.join(names), severity, summary, lat, lon]],
-                           columns=["Objects","Severity","Summary","Latitude","Longitude"])
-    history_df = pd.concat([history_df, new_row], ignore_index=True)
+        summary = incident_text
 
-    # Annotated image
+    # Annotated Image
     annotated = Image.fromarray(results[0].plot())
 
     # PDF
-    pdf_temp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
-    c = canvas.Canvas(pdf_temp.name)
-    c.drawString(80, 800, "InfraGuard AI Incident Report")
-    c.drawString(80, 770, f"Objects: {', '.join(names)}")
-    c.drawString(80, 750, f"Severity: {severity}")
-    c.drawString(80, 730, f"Summary: {summary}")
-    c.drawString(80, 710, f"Location: {lat}, {lon}")
-    c.save()
+    pdf_path = create_pdf(objects_text, severity, summary, lat, lon)
 
     # Map
-    html_map = generate_map(lat, lon)
-
-    return annotated, f"### Summary\n{summary}", severity_badge(severity), pdf_temp.name, html_map, history_df
-
-
-# ---------------- UI DESIGN ----------------
-with gr.Blocks(css=""" 
-    .gradio-container {max-width: 1100px !important; margin:auto;}
-    .header {text-align:center; font-size:40px; font-weight:800; margin-bottom:20px; color:#1E90FF;}
-    .card {padding:15px; border-radius:12px; border:1px solid #444; background:#111; margin-top:10px;}
-""") as demo:
-
-    gr.HTML("<div class='header'>🔥 InfraGuard — AI Security Incident Dashboard</div>")
-
-    with gr.Tabs():
-        # TAB 1
-        with gr.Tab("⚡ Analyze Incident"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    img_input = gr.Image(type="pil", label="Upload Image", height=350)
-                    lat = gr.Number(label="Latitude", value=24.8607)
-                    lon = gr.Number(label="Longitude", value=67.0011)
-                    analyze = gr.Button("🚀 Analyze", variant="primary")
-
-                with gr.Column(scale=2):
-                    output_img = gr.Image(label="Annotated Detection", height=350)
-                    output_summary = gr.Markdown()
-                    output_severity = gr.HTML()
-                    pdf = gr.File(label="📄 Download PDF Report")
-                    map_html = gr.HTML()
-
-        # TAB 2
-        with gr.Tab("📜 Incident History"):
-            history_table = gr.Dataframe(headers=["Objects","Severity","Summary","Latitude","Longitude"],
-                                         row_count=5, wrap=True)
-
-        # TAB 3
-        with gr.Tab("ℹ About"):
-            gr.Markdown("""
-            ### InfraGuard — AI-Powered Infrastructure Monitoring  
-            - Detects objects using YOLOv5  
-            - Generates professional incident summaries  
-            - Severity auto-classification  
-            - PDF reporting  
-            - Geolocation mapping  
-            """)
-
-    analyze.click(process_image,
-                  [img_input, lat, lon],
-                  [output_img, output_summary, output_severity, pdf, map_html, history_table])
-
-demo.launch()
+    map_html = generate_map(lat, lon)
+
+    # Save to history
+    thumb = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    img.save(thumb.name)
+    history.append({
+        "Image": thumb.name,
+        "Objects": objects_text,
+        "Severity": severity,
+        "Summary": summary,
+        "Latitude": lat,
+        "Longitude": lon
+    })
+
+    # History Table
+    hist_df = pd.DataFrame(history)
+
+    # Timeline Card
+    card_html = f"""
+    <div style='background:#0a1a2b;padding:15px;margin:10px;border-radius:10px;border:1px solid #0ff;'>
+        <h3 style='color:#00eaff'>New Incident Logged</h3>
+        <p><b>Objects:</b> {objects_text}</p>
+        <p><b>Severity:</b> {severity_badge(severity)}</p>
+        <p><b>Summary:</b> {summary}</p>
+        <p><b>Location:</b> {lat}, {lon}</p>
+    </div>
+    """
+
+    return annotated, summary, severity, pdf_path, map_html, card_html, hist_df
+
+
+# ------------------------------
+# CUSTOM CSS (CYBER BLUE THEME)
+# ------------------------------
+cyber_css = """
+body { background-color: #051622 !important; }
+
+.gradio-container {
+    background: #051622 !important;
+    color: white !important;
+}
+
+h1 {
+    color: #00eaff !important;
+    text-align:center;
+    font-family: 'Arial Black';
+}
+
+.gr-button {
+    background:#00eaff !important;
+    color:black !important;
+    border-radius:8px;
+    font-weight:bold;
+    border:none !important;
+}
+"""
+
+# ------------------------------
+# UI LAYOUT
+# ------------------------------
+with gr.Blocks(css=cyber_css, theme=gr.themes.Soft()) as demo:
+    gr.HTML("<h1>🔥 INFRA GUARD – Cyber Security AI Dashboard</h1>")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_in = gr.Image(type="pil", label="Upload Scene")
+            lat = gr.Number(value=24.8607, label="Latitude")
+            lon = gr.Number(value=67.0011, label="Longitude")
+            btn = gr.Button("Analyze Incident")
+        with gr.Column(scale=1):
+            img_out = gr.Image(label="Annotated Image")
+            summary = gr.Textbox(label="AI Summary")
+            severity = gr.Textbox(label="Severity")
+            pdf = gr.File(label="Download PDF Report")
+            map_view = gr.HTML(label="Incident Map")
+
+    gr.Markdown("### 🕒 Timeline Feed")
+    timeline = gr.HTML()
+
+    gr.Markdown("### 📘 Incident History")
+    hist_table = gr.Dataframe(headers=["Image", "Objects", "Severity", "Summary", "Latitude", "Longitude"], value=[])
+
+    btn.click(
+        process,
+        inputs=[img_in, lat, lon],
+        outputs=[img_out, summary, severity, pdf, map_view, timeline, hist_table]
+    )
+
+demo.launch()