|
| 1 | +import streamlit as st |
| 2 | +import cv2 |
| 3 | +import numpy as np |
| 4 | +import onnxruntime as ort |
| 5 | +from PIL import Image |
| 6 | +import time |
| 7 | +import os |
| 8 | +from datetime import datetime |
| 9 | + |
| 10 | +# ========================= |
| 11 | +# GLOBALS & SETTINGS |
| 12 | +# ========================= |
| 13 | +ort.set_default_logger_severity(3) |
| 14 | + |
| 15 | +CLASS_NAMES = { |
| 16 | + 0: "good", |
| 17 | + 1: "broken", |
| 18 | + 2: "flashover" |
| 19 | +} |
| 20 | + |
| 21 | +BBOX_COLOR = (0, 255, 0) # Bright Green |
| 22 | +BBOX_THICKNESS = 3 |
| 23 | + |
| 24 | +# Ensure output directory exists locally |
| 25 | +OUTPUT_DIR = "output" |
| 26 | +os.makedirs(OUTPUT_DIR, exist_ok=True) |
| 27 | + |
| 28 | +# ========================= |
| 29 | +# CACHED MODEL LOADING |
| 30 | +# ========================= |
| 31 | +@st.cache_resource |
| 32 | +def load_model(): |
| 33 | + model_path = 'best.onnx' |
| 34 | + |
| 35 | + providers = [ |
| 36 | + ( |
| 37 | + "TensorrtExecutionProvider", |
| 38 | + { |
| 39 | + "trt_fp16_enable": True, |
| 40 | + "trt_engine_cache_enable": True, |
| 41 | + "trt_engine_cache_path": "./trt_cache", |
| 42 | + "trt_max_workspace_size": 1073741824 |
| 43 | + } |
| 44 | + ), |
| 45 | + "CUDAExecutionProvider" |
| 46 | + ] |
| 47 | + |
| 48 | + sess_options = ort.SessionOptions() |
| 49 | + sess_options.log_severity_level = 3 |
| 50 | + |
| 51 | + session = ort.InferenceSession( |
| 52 | + model_path, |
| 53 | + sess_options=sess_options, |
| 54 | + providers=providers |
| 55 | + ) |
| 56 | + |
| 57 | + input_name = session.get_inputs()[0].name |
| 58 | + return session, input_name |
| 59 | + |
| 60 | +# ========================= |
| 61 | +# INFERENCE FUNCTION |
| 62 | +# ========================= |
| 63 | +def run_inference(img, session, input_name, conf_thresh): |
| 64 | + img_h, img_w = img.shape[:2] |
| 65 | + |
| 66 | + # Preprocess |
| 67 | + input_img = cv2.resize(img, (640, 640)) |
| 68 | + input_img = input_img.transpose(2, 0, 1) |
| 69 | + input_img = np.expand_dims(input_img, axis=0).astype(np.float32) / 255.0 |
| 70 | + |
| 71 | + # Inference |
| 72 | + start_time = time.time() |
| 73 | + outputs = session.run(None, {input_name: input_img}) |
| 74 | + infer_time = (time.time() - start_time) * 1000 |
| 75 | + |
| 76 | + # Postprocess |
| 77 | + predictions = np.squeeze(outputs[0]).T |
| 78 | + boxes, scores, class_ids = [], [], [] |
| 79 | + |
| 80 | + x_factor = img_w / 640 |
| 81 | + y_factor = img_h / 640 |
| 82 | + |
| 83 | + for row in predictions: |
| 84 | + class_scores = row[4:] |
| 85 | + max_score = np.max(class_scores) |
| 86 | + |
| 87 | + if max_score > conf_thresh: |
| 88 | + class_id = np.argmax(class_scores) |
| 89 | + xc, yc, w, h = row[0], row[1], row[2], row[3] |
| 90 | + |
| 91 | + x_min = int((xc - w / 2) * x_factor) |
| 92 | + y_min = int((yc - h / 2) * y_factor) |
| 93 | + width = int(w * x_factor) |
| 94 | + height = int(h * y_factor) |
| 95 | + |
| 96 | + boxes.append([x_min, y_min, width, height]) |
| 97 | + scores.append(float(max_score)) |
| 98 | + class_ids.append(class_id) |
| 99 | + |
| 100 | + indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold=conf_thresh, nms_threshold=0.4) |
| 101 | + |
| 102 | + defects_found = len(indices) if len(indices) > 0 else 0 |
| 103 | + |
| 104 | + # Draw Boxes |
| 105 | + if defects_found > 0: |
| 106 | + for i in indices.flatten(): |
| 107 | + x, y, w, h = boxes[i] |
| 108 | + current_class_id = class_ids[i] |
| 109 | + class_name = CLASS_NAMES.get(current_class_id, f"Class {current_class_id}") |
| 110 | + |
| 111 | + cv2.rectangle(img, (x, y), (x + w, y + h), BBOX_COLOR, BBOX_THICKNESS) |
| 112 | + label = f"{class_name}: {scores[i]:.2f}" |
| 113 | + (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) |
| 114 | + cv2.rectangle(img, (x, y - th - 10), (x + tw, y), BBOX_COLOR, -1) |
| 115 | + cv2.putText(img, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2) |
| 116 | + |
| 117 | + return img, defects_found, infer_time |
| 118 | + |
| 119 | +# ========================= |
| 120 | +# STREAMLIT UI LAYOUT |
| 121 | +# ========================= |
| 122 | +st.set_page_config(page_title="Defect Diagnosis UI", layout="wide") |
| 123 | + |
| 124 | +st.title("🔍 Edge AI Defect Diagnosis") |
| 125 | +st.markdown("Run TensorRT-optimized YOLOv8 inference directly from your browser.") |
| 126 | + |
| 127 | +# Load Model |
| 128 | +with st.spinner('Loading TensorRT Engine...'): |
| 129 | + session, input_name = load_model() |
| 130 | + |
| 131 | +# Sidebar controls |
| 132 | +st.sidebar.header("Settings") |
| 133 | +confidence_threshold = st.sidebar.slider( |
| 134 | + "Confidence Threshold", |
| 135 | + min_value=0.1, max_value=1.0, value=0.3, step=0.05 |
| 136 | +) |
| 137 | + |
| 138 | +st.sidebar.divider() |
| 139 | + |
| 140 | +# Input Method Selector |
| 141 | +input_method = st.sidebar.radio( |
| 142 | + "Choose Input Method:", |
| 143 | + ("📁 Upload Image(s)", "📷 Webcam Capture") |
| 144 | +) |
| 145 | + |
| 146 | +st.sidebar.info(f"Processed images are automatically saved to `./{OUTPUT_DIR}/`") |
| 147 | + |
| 148 | +# ---------------------------------------- |
| 149 | +# MODE 1: FILE UPLOAD (Single or Multiple) |
| 150 | +# ---------------------------------------- |
| 151 | +if input_method == "📁 Upload Image(s)": |
| 152 | + uploaded_files = st.file_uploader( |
| 153 | + "Upload images (Select multiple files to process a batch)", |
| 154 | + type=["jpg", "jpeg", "png"], |
| 155 | + accept_multiple_files=True |
| 156 | + ) |
| 157 | + |
| 158 | + if uploaded_files: |
| 159 | + cols = st.columns(2) # Create a 2-column grid |
| 160 | + |
| 161 | + for idx, uploaded_file in enumerate(uploaded_files): |
| 162 | + # Read image |
| 163 | + image = Image.open(uploaded_file) |
| 164 | + img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) |
| 165 | + |
| 166 | + # Run Inference |
| 167 | + result_img, defect_count, inf_time = run_inference( |
| 168 | + img_array, session, input_name, confidence_threshold |
| 169 | + ) |
| 170 | + |
| 171 | + # Save to output folder |
| 172 | + output_filepath = os.path.join(OUTPUT_DIR, uploaded_file.name) |
| 173 | + cv2.imwrite(output_filepath, result_img) |
| 174 | + |
| 175 | + # Display in UI |
| 176 | + result_img_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB) |
| 177 | + with cols[idx % 2]: |
| 178 | + st.image(result_img_rgb, caption=f"{uploaded_file.name}") |
| 179 | + st.write(f"**Defects Found:** {defect_count} | **Time:** {inf_time:.2f} ms") |
| 180 | + st.divider() |
| 181 | + |
| 182 | +# ---------------------------------------- |
| 183 | +# MODE 2: WEBCAM CAPTURE |
| 184 | +# ---------------------------------------- |
| 185 | +elif input_method == "📷 Webcam Capture": |
| 186 | + # Streamlit native webcam widget |
| 187 | + camera_image = st.camera_input("Take a picture to analyze") |
| 188 | + |
| 189 | + if camera_image is not None: |
| 190 | + # Read image from camera buffer |
| 191 | + image = Image.open(camera_image) |
| 192 | + img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) |
| 193 | + |
| 194 | + # Run Inference |
| 195 | + result_img, defect_count, inf_time = run_inference( |
| 196 | + img_array, session, input_name, confidence_threshold |
| 197 | + ) |
| 198 | + |
| 199 | + # Generate a unique filename using timestamp |
| 200 | + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 201 | + filename = f"webcam_{timestamp}.jpg" |
| 202 | + output_filepath = os.path.join(OUTPUT_DIR, filename) |
| 203 | + |
| 204 | + # Save to output folder |
| 205 | + cv2.imwrite(output_filepath, result_img) |
| 206 | + |
| 207 | + # Display in UI |
| 208 | + result_img_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB) |
| 209 | + |
| 210 | + st.success(f"Image processed and saved locally as `{filename}`") |
| 211 | + st.image(result_img_rgb, caption="Webcam Analysis Result") |
| 212 | + st.write(f"**Defects Found:** {defect_count}") |
| 213 | + st.write(f"**Inference Time:** {inf_time:.2f} ms") |
0 commit comments