app.py

import streamlit as st
import numpy as np
import cv2
import tempfile
import time
from detection import AvgPoint, detect_catenary_lines, detect_pantograph_line, point, line_points, intersection

st.set_page_config(layout="wide")
#video_file = st.file_uploader("Загрузите видео")
f = st.file_uploader("Choose a Video") 

if f:
    tfile = tempfile.NamedTemporaryFile(delete=False) 
    tfile.write(f.read())

    # Opens the Video file
    cap= cv2.VideoCapture(tfile.name)
    ret, frame = cap.read()


    col1, col2 = st.columns(2)
    with col1:
        pant_left = st.slider("Введите координату левого края пантографа, pt:", min_value = 400, max_value = 500, value = 420)
        pant_right = st.slider("Введите координату правого края пантографа, pt:", min_value = 720, max_value = 800)
        pant_len = st.number_input("Введите длину пантографа, мм:", min_value = 1000)
        thresh_mm     = st.number_input("Введите величину максимального отклонения от центра, мм:", min_value = 300)
        pant_center = (int((pant_right - pant_left)/2 + pant_left), 350)      
        mm_px = (pant_right - pant_left)/pant_len
        thresh_px = thresh_mm * mm_px
        thresh_left = pant_center[0] - int(thresh_px)
        thresh_right = pant_center[0] + int(thresh_px)
        pant_top = 310
        pant_bottom = 390
    with col2:
        
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        cv2.line(rgb_frame, (pant_left, pant_top), (pant_left, pant_bottom), [0, 255, 0], 3)    
        cv2.line(rgb_frame, (pant_right, pant_top), (pant_right, pant_bottom), [0, 255, 0], 3)   
  
        cv2.line(rgb_frame, (thresh_left, pant_top+20), (thresh_left, pant_bottom-20), [255, 0, 0 ], 3)    
        cv2.line(rgb_frame, (thresh_right, pant_top+20), (thresh_right, pant_bottom-20), [255, 0, 0 ], 3)   
        
        dot_color = [255, 255, 255]
        dot_size = 6
        cv2.circle(rgb_frame, pant_center, dot_size, dot_color, -1)
        st.image(rgb_frame)

    start_video = st.button("Старт видео")
    
    pant_borders = [pant_left, pant_right, pant_top, pant_bottom]
    pant_y_coord = int((pant_bottom - pant_top)/2 + pant_top)
    template = [pant_left, 350, pant_right, 350]
    template_norm = np.linalg.norm(template)
    gray  = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 40, 700)


    # get lines
    # (x1, y1, x2, y2)
    lines = cv2.HoughLinesP(
        edges,
        rho=1.5,
        theta=np.pi/180,
        threshold=50,
        minLineLength=80,
        maxLineGap=40        
    )
    # Pantograph start position

    pantograph = detect_pantograph_line(lines, pant_borders, template_norm)         
    cur_pantograph = pantograph
    frame_count = 0
    flag_anomaly = True
    anomaly_count = 0
    if start_video:
        stframe = st.empty()

        while cap.isOpened():
            #time.sleep(0.1)
            ret, frame = cap.read()
            # if frame is read correctly ret is True
            if not ret:
                print("Can't receive frame (stream end?). Exiting ...")
                break

            gray  = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 40, 700)


            # get lines
            # (x1, y1, x2, y2)
            lines = cv2.HoughLinesP(
                edges,
                rho=1.5,
                theta=np.pi/180,
                threshold=50,
                minLineLength=80,
                maxLineGap=40        
            )
            
#------------------------ Pantograph current position-----------------------
  
            
            pantograph = detect_pantograph_line(lines, pant_borders, template_norm)  
            mean_y = (pantograph[1] + pantograph[3])/2
            mean_y_old = (cur_pantograph[1] + cur_pantograph[3])/2
            if abs(mean_y - mean_y_old) > 20:
                pantograph = cur_pantograph
            else:
                cur_pantograph = pantograph
            
            cleaned_img = np.zeros((frame.shape[0], frame.shape[1], 3), dtype=np.uint8)
            pant_line_color = [0, 255, 0]
            pant_line_thickness = 3
            img = cv2.line(cleaned_img, (pant_left, pantograph[1]), (pant_right, pantograph[3]), pant_line_color, pant_line_thickness)

#---------------------- Patnograph center-------------------------

            pant_center_x = (pant_right - pant_left)/2 + pant_left
            pant_center_y = (pantograph[3] - pantograph[1])/2 + pantograph[1]
#            avg_pant = AvgPoint()
#            avg_pant.append(pant_center_x, pant_center_y)
#            x = int(avg_pant.avg_x)
#            y = int(avg_pant.avg_y)
            x = int(pant_center_x)
            y = int(pant_center_y)
            pant_center_point = (x, y)
            dot_color = [255, 255, 255]
            dot_size = 6

            img = cv2.circle(cleaned_img, pant_center_point, dot_size, dot_color, -1)

            #Catenary lines
            catenary_lines = detect_catenary_lines(lines)
            line_color = [0, 0, 255]
            line_thickness = 1
            for line in catenary_lines:
                for x1, y1, x2, y2 in line:
                    img = cv2.line(cleaned_img, (x1, y1), (x2, y2), line_color, line_thickness)


            #Pantogrph/catenary contact point
            if len(catenary_lines) > 0:
                contact = intersection(np.array(pantograph).reshape((1, 4)), catenary_lines[0])
                
            if contact:
                dot_color = [255, 0, 255]
                dot_size = 6
                img = cv2.circle(cleaned_img, contact, dot_size, dot_color, -1)

                #Point relative 
                dist = np.linalg.norm(np.array(pant_center_point)- np.array(contact)).astype(int)


                if dist > thresh_px:
                    if flag_anomaly:
                        anomaly_count += 1
                    flag_anomaly = False
                else:
                    flag_anomaly = True 

#--------------------- Message-----------------------------------------
            font = cv2.FONT_HERSHEY_COMPLEX
            if frame_count % 50 == 0:    
                print_dist = min(dist, thresh_mm)
            frame_count += 1
            img = cv2.putText(cleaned_img, "Отклонение:" + str(print_dist) + "mm", (10,800), font, 1, (0,255,0))
                
            img = cv2.putText(cleaned_img, "Количество отклонений:" + str(anomaly_count), (10, 850), font, 1, (0,255,0))        
        
#            font = cv2.FONT_HERSHEY_COMPLEX
#            img = cv2.putText(cleaned_img, "Отклонение:" + str(mean_y) + "mm",(10,800), font, 1, (0,255,0))
#            img = cv2.putText(cleaned_img, "Отклонение:" + str(mean_y_old) + "mm",(10,900), font, 1, (0,255,0))
            
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  
            overlay = cv2.addWeighted(frame, 0.8, cleaned_img, 1.0, 0.0)
            stframe.image(overlay)