File size: 6,098 Bytes
e5fa36b |
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 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import streamlit as st
from PIL import Image
import matplotlib.pyplot as plt
import networkx as nx
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
torch.cuda.empty_cache()
#import openai
import os
import numpy as np
# Ensure models and datasets are available
from download_models import download_all
# Run only if critical file is missing
if not os.path.exists("P&ID-Symbols-3/train/_annotations.coco.json"):
with st.spinner("Downloading required files (models & datasets)..."):
download_all()
from pipeline.detector import detect_symbols_and_lines
from pipeline.graph_builder import build_graph
from pipeline.gnn_model import run_gnn
from pipeline.agent import generate_agent_actions
st.set_page_config(layout="wide")
st.title(" Agentic Predictive Maintenance (P&ID Graph + GNN)")
# Initialize session state variables
if "G" not in st.session_state:
st.session_state.G = None
if "feature_map" not in st.session_state:
st.session_state.feature_map = {}
if "scores" not in st.session_state:
st.session_state.scores = {}
#uploaded_file = st.file_uploader("Upload a P&ID Image", type=["png", "jpg", "jpeg"])
#if uploaded_file:
# === User can choose from downloaded dataset OR upload their own ===
st.subheader("Upload or Select a P&ID Drawing")
local_dataset_dir = "P&ID-Symbols-3/P&ID-Symbols-3/test"
image_files = []
if os.path.exists(local_dataset_dir):
image_files = [f for f in os.listdir(local_dataset_dir) if f.lower().endswith((".png", ".jpg", ".jpeg"))]
else:
st.warning(f"Dataset folder not found: {local_dataset_dir}. Please run download_models.py to download it.")
selected_image = st.selectbox("Select a sample from P&ID-Symbols-3:", ["-- Select an example --"] + image_files)
uploaded_file = st.file_uploader("...Or upload your own P&ID image", type=["png", "jpg", "jpeg"])
image = None
image_source = ""
if selected_image and selected_image != "-- Select an example --":
image_path = os.path.join(local_dataset_dir, selected_image)
image = Image.open(image_path)
image_source = f"Sample from dataset: {selected_image}"
elif uploaded_file:
image = Image.open(uploaded_file)
image_source = f"Uploaded: {uploaded_file.name}"
if image:
st.image(image, caption=image_source, use_column_width=True)
#image = Image.open(uploaded_file)
#st.image(image, caption="P&ID Diagram", use_column_width=True)
if st.button(" Run Detection and Analysis"):
# Uncomment these when detection and graph building pipelines are ready
detections, annotations, class_names = detect_symbols_and_lines(image)
graph = build_graph(image, detections, annotations, class_names)
st.info("Running anomaly detection on the graph (simulated for now)...")
fig, feature_map, red_nodes, central_node, scores, G = run_gnn()
st.session_state.G = G
st.session_state.feature_map = feature_map
st.session_state.scores = scores
st.pyplot(fig)
actions = generate_agent_actions(fig, feature_map, red_nodes, central_node, scores)
for action in actions:
st.write(action)
# === DeepSeek Local Model Setup ===
@st.cache_resource
def load_deepseek_model():
model_name = "deepseek-ai/deepseek-coder-1.3b-instruct" # Lightweight version
# model_name = "deepseek-ai/deepseek-llm-7b" # Larger but more capable
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
'''model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True'''
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="cpu",
#load_in_4bit=True, # 4-bit quantization
trust_remote_code=True
)
return model, tokenizer
# === Q&A Interface ===
st.subheader(" Ask Questions About the Graph (DeepSeek Local)")
user_query = st.chat_input("Ask a question about the graph...")
if user_query:
G = st.session_state.get("G")
feature_map = st.session_state.get("feature_map", {})
scores = st.session_state.get("scores", [])
if G is not None and feature_map and len(scores) > 0:
graph_data = {
"nodes": [
{
"id": str(i),
"label": feature_map[i] if i < len(feature_map) else f"Node {i}",
"score": float(scores[i]) if i < len(scores) else 0.0
}
for i in G.nodes()
],
"edges": [
{"source": str(u), "target": str(v)}
for u, v in G.edges()
]
}
prompt = (
"You are an expert graph analyst. Analyze this P&ID graph and answer the question.\n\n"
"### Graph Data:\n"
f"{json.dumps(graph_data, indent=2)}\n\n"
"### Question:\n"
f"{user_query}\n\n"
"### Answer:\n"
)
try:
with st.spinner("Thinking (via DeepSeek Local)..."):
# Load model (cached after first run)
model, tokenizer = load_deepseek_model()
# Generate response
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=128,
temperature=0.7,
do_sample=True
)
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Remove the prompt from the answer
answer = answer[len(prompt):].strip()
st.markdown(f"**DeepSeek:** {answer}")
except Exception as e:
st.error(f"DeepSeek error: {e}")
st.error("Make sure you have enough GPU memory (8GB+ recommended for 7B model)")
else:
st.warning("Graph or scores are not ready yet.")
|