# ----------------------------------------------------------------------------- # # This file is part of the PantoScanner distribution on: # https://huggingface.co/spaces/swissrail/PantoScanner # # PantoScanner - Analytics and measurement capability for technical objects. # Copyright (C) 2017-2024 Schweizerische Bundesbahnen SBB # # Authors (C) 2024 L. Hofstetter (lukas.hofstetter@sbb.ch) # Authors (C) 2017 U. Gehrig (urs.gehrig@sbb.ch) # Authors (C) 2026 U. Gehrig (urs@lighthouse.no) # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . # # ----------------------------------------------------------------------------- import streamlit as st import os import glob import cv2 import numpy as np from strip_measure_4_0 import prepare_networks_for_measurement, measure_strip import altair as alt import pandas as pd IMG_BASE_DIR = 'images' CAMERA_MATRIX = [ [11100, 0, 1604], [0, 11100, 1100], [0, 0, 1] ] OBJECT_REFERENCE_POINTS = [ [347, 0, 42], # B [347, 0, 522], # D [-347, 26, 480], # F [-347, 26, 0]] # H LOWER_CONTOUR_QUADRATIC_CONSTANT = 0.00005 CAMERA_PARAMETERS = (50, 0.0045, 2200, 3208) PLANE_PARAMETERS_CLOSE = ([0, 0, 0], (1, 0, 0), (0, 1, 0)) PLANE_PARAMETERS_FAR = ([0, 0, 480], (1, 0, 0), (0, 1, 0)) BOUNDARY_1 = (300, 92) BOUNDARY_2 = (650, 1500) IMAGE_SIZE_SEG = 1408 IMAGE_WIDTH_SEG = 1408 IMAGE_HEIGHT_SEG = 576 path_yolo_model = os.path.join(os.getcwd(), 'app', 'detection_model.pt') path_segmentation_model = os.path.join(os.getcwd(), 'app', 'segmentation_model.pth') # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def get_image_paths(base_dir: str): pattern = os.path.join(os.getcwd(), base_dir, '*.png') return sorted(glob.glob(pattern)) def get_num_images(): return len(st.session_state['image_path_list']) def increment_index(index_current, max_index, overflow=False, min_index=0): index_new = index_current + 1 if index_new <= max_index: return index_new return min_index if overflow else index_current def decrement_index(index_current, min_index, overflow=False, max_index=-1): index_new = index_current - 1 if index_new >= min_index: return index_new return max_index if overflow else index_current def load_image_array(image_path: str) -> np.ndarray: bgr = cv2.imread(image_path) if bgr is None: raise FileNotFoundError(f"Cannot read image: {image_path}") return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) def get_current_image() -> np.ndarray: idx = st.session_state['image_index_current'] cached = st.session_state['image_data_list'][idx] if isinstance(cached, np.ndarray): return cached img = load_image_array(st.session_state['image_path_list'][idx]) st.session_state['image_data_list'][idx] = img return img def get_current_measurement(): idx = st.session_state['image_index_current'] m = st.session_state['measurement_data_list'][idx] return (True, m) if m is not None else (False, None) def update_measurements(measurement, index_measurement: int): st.session_state['measurement_data_list'][index_measurement] = measurement # --------------------------------------------------------------------------- # Callbacks # --------------------------------------------------------------------------- def update_on_index_change(new_index: int): st.session_state['image_index_current'] = new_index st.session_state['current_image_array'] = get_current_image() st.session_state['current_measurement'] = get_current_measurement() def callback_button_previous(overflow_index: bool = True): new_index = decrement_index( st.session_state['image_index_current'], min_index=0, overflow=overflow_index, max_index=st.session_state['num_images'] - 1 ) update_on_index_change(new_index) def callback_button_next(overflow_index: bool = True): new_index = increment_index( st.session_state['image_index_current'], st.session_state['num_images'], overflow=overflow_index, min_index=0 ) update_on_index_change(new_index) # --------------------------------------------------------------------------- # Measurement + display # --------------------------------------------------------------------------- def measure_image(image_path: str): result = measure_strip( img_path=image_path, model_yolo=st.session_state['models']['detection'], segmentation_model=st.session_state['models']['segmentation'], camera_matrix=CAMERA_MATRIX, object_reference_points=OBJECT_REFERENCE_POINTS, camera_parameters=CAMERA_PARAMETERS, plane_parameters_close=PLANE_PARAMETERS_CLOSE, plane_parameters_far=PLANE_PARAMETERS_FAR, lower_contour_quadratic_constant=LOWER_CONTOUR_QUADRATIC_CONSTANT, boundary_1=BOUNDARY_1, boundary_2=BOUNDARY_2, image_size_seg=IMAGE_SIZE_SEG, image_width_seg=IMAGE_WIDTH_SEG, image_height_seg=IMAGE_HEIGHT_SEG, ) arr_0, arr_1 = result[0], result[1] arr_0[:, 0] = np.abs(arr_0[:, 0]) arr_1[:, 0] = np.abs(arr_1[:, 0]) return arr_0, arr_1 def build_table(profile_a: np.ndarray, profile_b: np.ndarray) -> pd.DataFrame: """Merge both profiles into a single table keyed on position.""" df_a = pd.DataFrame({'Position': profile_a[:, 1], 'Profile A (mm)': profile_a[:, 0]}) df_b = pd.DataFrame({'Position': profile_b[:, 1], 'Profile B (mm)': profile_b[:, 0]}) df = (pd.merge(df_a, df_b, on='Position', how='outer') .sort_values('Position') .reset_index(drop=True) .round(2)) return df def build_chart(profile_a: np.ndarray, profile_b: np.ndarray): """Build Altair chart from two profile arrays (col0=height, col1=position).""" rows = [] for h, p in zip(profile_a[:, 0].tolist(), profile_a[:, 1].tolist()): rows.append({'Position (mm)': int(p), 'Height (mm)': round(float(h), 3), 'Strip': 'Profile A'}) for h, p in zip(profile_b[:, 0].tolist(), profile_b[:, 1].tolist()): rows.append({'Position (mm)': int(p), 'Height (mm)': round(float(h), 3), 'Strip': 'Profile B'}) df_chart = pd.DataFrame(rows).sort_values('Position (mm)').reset_index(drop=True) chart = ( alt.Chart(df_chart, title='Carbon strip profile — height across width') .mark_line(point=False, strokeWidth=1.5) .encode( x=alt.X('Position (mm):Q', title='Position (mm)'), y=alt.Y('Height (mm):Q', title='Height (mm)', scale=alt.Scale(zero=False)), color=alt.Color( 'Strip:N', scale=alt.Scale(domain=['Profile A', 'Profile B'], range=['#1f77b4', '#d62728']), legend=alt.Legend(title='Strip'), ), tooltip=['Position (mm):Q', 'Height (mm):Q', 'Strip:N'], ) .properties(height=300) .interactive() ) return chart def render_right_panel(): """Render table + chart — shows placeholder if no measurement yet.""" has_measurement, measurement_data = get_current_measurement() if not has_measurement: st.markdown( '
' '📏 Click  Measure  to analyse this image' '
', unsafe_allow_html=True, ) return profile_a = np.array(measurement_data[0], dtype=float) profile_b = np.array(measurement_data[1], dtype=float) df = build_table(profile_a, profile_b) st.subheader('Profile heights (mm)') st.dataframe( df.style.format({'Position': '{:.1f}', 'Profile A (mm)': '{:.2f}', 'Profile B (mm)': '{:.2f}'}), use_container_width=True, height=380, ) fig = build_chart(profile_a, profile_b) st.altair_chart(fig, use_container_width=True) def callback_button_measure(): has_measurement, _ = get_current_measurement() if not has_measurement: with st.spinner('Running AI measurement…'): img_path = st.session_state['image_path_list'][ st.session_state['image_index_current'] ] result = measure_image(img_path) update_measurements(result, st.session_state['image_index_current']) st.session_state['current_measurement'] = get_current_measurement() # --------------------------------------------------------------------------- # Session-state initialisation # --------------------------------------------------------------------------- def _init_session_state(): if 'image_path_list' not in st.session_state: st.session_state['image_path_list'] = get_image_paths(IMG_BASE_DIR) if 'num_images' not in st.session_state: st.session_state['num_images'] = get_num_images() if 'image_data_list' not in st.session_state: st.session_state['image_data_list'] = [None] * st.session_state['num_images'] if 'image_index_current' not in st.session_state: st.session_state['image_index_current'] = 0 if 'current_image_array' not in st.session_state: st.session_state['current_image_array'] = get_current_image() if 'measurement_data_list' not in st.session_state: st.session_state['measurement_data_list'] = [None] * st.session_state['num_images'] if 'current_measurement' not in st.session_state: st.session_state['current_measurement'] = get_current_measurement() if 'models' not in st.session_state: with st.spinner('Loading AI models…'): seg_model, yolo_model = prepare_networks_for_measurement( model_yolo_path=path_yolo_model, model_segmentation_path=path_segmentation_model, ) st.session_state['models'] = {'segmentation': seg_model, 'detection': yolo_model} # --------------------------------------------------------------------------- # Page layout # --------------------------------------------------------------------------- st.set_page_config(layout='wide', page_title='PantoScanner') _init_session_state() st.title('🚃 PantoScanner') st.caption('Detects and segments a pantograph carbon strip, then extracts thickness as a height profile.') if st.session_state['num_images'] == 0: st.error('No images found in the `images/` directory.') st.stop() # Two-column layout: image + controls on left, results on right left_col, right_col = st.columns([2, 3]) with left_col: st.image( st.session_state['current_image_array'], caption=f"Image {st.session_state['image_index_current'] + 1} of {st.session_state['num_images']}", use_column_width=True, ) b1, b2, b3 = st.columns(3) with b1: st.button('⬅ Prev', on_click=callback_button_previous, kwargs={'overflow_index': True}, use_container_width=True) with b2: measure_clicked = st.button('📏 Measure', use_container_width=True, type='primary') with b3: st.button('Next ➡', on_click=callback_button_next, kwargs={'overflow_index': True}, use_container_width=True) if measure_clicked: callback_button_measure() st.success('Measurement done!') elif get_current_measurement()[0]: st.success('Measurement done!') with right_col: render_right_panel()