Spaces:

swissrail
/

PantoScanner

Runtime error

Update app.py

616b22b verified 7 months ago

12.3 kB

# -----------------------------------------------------------------------------
#
# 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 <https://www.gnu.org/licenses/>.
#
# -----------------------------------------------------------------------------

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(
            '<div style="height:200px;display:flex;align-items:center;justify-content:center;'
            'border:1px dashed #ccc;border-radius:8px;color:#999;font-size:1.1rem;">'
            '📏 Click <b style="color:#555">&nbsp;Measure&nbsp;</b> to analyse this image'
            '</div>',
            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()

	# -----------------------------------------------------------------------------
	#
	# 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 <https://www.gnu.org/licenses/>.
	#
	# -----------------------------------------------------------------------------

	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(
	'<div style="height:200px;display:flex;align-items:center;justify-content:center;'
	'border:1px dashed #ccc;border-radius:8px;color:#999;font-size:1.1rem;">'
	'📏 Click <b style="color:#555"> Measure </b> to analyse this image'
	'</div>',
	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()