# -*- coding: utf-8 -*- """ Luke Fullard: 16 June 2024 Script to track an object of interest in a video """ import streamlit as st import cv2 import numpy as np from PIL import Image import tempfile import pandas as pd import os from io import BytesIO import base64 from skimage import filters ############################################################################### ############################################################################### ############################################################################### # Function to apply image adjustments with new options def apply_adjustments(frame, grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold): if grayscale: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if contrast > 1.0: # Clip the values to avoid wrapping around frame = np.clip(contrast * frame, 0, 255).astype(np.uint8) if blur > 0: frame = cv2.GaussianBlur(frame, (blur, blur), 0) if edges: frame = cv2.Canny(frame, lower_threshold, upper_threshold) # Sharpening if sharpen_amount > 0: frame = filters.unsharp_mask(frame, radius=1, amount=sharpen_amount) frame = (frame * 255).astype(np.uint8) # Convert back to uint8 # Background Subtraction if subtract_background: # Initialize the background subtractor fgbg = cv2.createBackgroundSubtractorMOG2() frame = fgbg.apply(frame) if binarize: _, frame = cv2.threshold(frame, binarize_threshold, 255, cv2.THRESH_BINARY) return frame ############################################################################### ############################################################################### ############################################################################### def part_one(): st.subheader('Part one: upload your video file') with st.expander('Instructions #1'): st.write('Please choose a file to upload. At present only the following formats are supported: "mp4", "avi", "mov", "mkv"') uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "avi", "mov", "mkv"]) return uploaded_file ############################################################################### ############################################################################### ############################################################################### def part_two(uploaded_file): st.subheader('Part two: adjust image and select region of interest') with st.expander('Instructions #2'): st.write(''' Part two is a two step process: a) Use the Image adjustments in the left hand side toolbar to modify the image until suitable for tracking analysis. In practice, I have found the "Binarize" tool to be the most useful for tracking. You can use the frame number selector just above the image to see how the image adjustment affects the other frames in the video sequence. b) Set the region of interest (ROI). Using the four number sliders below, choose the initial area of interest. the X,Y numbers define the top left of the rectangular area of interest, while the width and height define the rectangle dimensions. The rectangle will be drawn on the image to help guide you. **NOTE: The ROI is always set on the first frame, so ensure you are on that image frame when setting the ROI.** Once you are ready, click the "Start Object Tracking" button below the image. ''') # Save the uploaded video file to a temporary file tfile = tempfile.NamedTemporaryFile(delete=False) tfile.write(uploaded_file.read()) file_name, _ = os.path.splitext(uploaded_file.name) # Read the first frame of the video from the temporary file vcap = cv2.VideoCapture(tfile.name) total_frames = int(vcap.get(cv2.CAP_PROP_FRAME_COUNT)) # Close and delete the temporary file tfile.close() instructions_container = st.empty() left_column,right_column = st.columns(2) # User input for frame number frame_number = st.number_input(f'Enter a frame number between 1 and {total_frames}', min_value=1, max_value=total_frames, value=1) if frame_number >1: st.warning("WARNING: When setting the region of interest (ROI) please ensure that you are defining the region based on Frame #1.") # Options for adjustments st.sidebar.header("Image adjustments") grayscale = st.sidebar.checkbox("Convert to Grayscale") contrast = st.sidebar.slider("Contrast", 1.0, 3.0, 1.0) blur = st.sidebar.slider("Blur", 0, 10, 0) edges = st.sidebar.checkbox("Edge Detection") # Options for edge detection adjustments lower_threshold = 0 upper_threshold = 0 if edges: st.sidebar.header("Edge Detection Parameters") lower_threshold = st.sidebar.slider("Lower Threshold", 0, 255, 100) upper_threshold = st.sidebar.slider("Upper Threshold", 0, 255, 200) sharpen_amount = st.sidebar.slider("Sharpen Amount", 0.0, 2.0, 0.0) subtract_background = st.sidebar.checkbox("Background Subtraction") # Options for binarization adjustments binarize = st.sidebar.checkbox("Binarize") binarize_threshold=0 if binarize: st.sidebar.header("Binarization Parameters") binarize_threshold = st.sidebar.slider("Binarization Threshold", 0, 255, 128) # Seek to the specified frame vcap.set(cv2.CAP_PROP_POS_FRAMES, frame_number-1) success, frame = vcap.read() if success: # Manual ROI entry with instructions_container: st.write("Enter the Region of interest (ROI) coordinates") with left_column: x = st.number_input('Enter X coordinate of top-left corner', min_value=0, value=0) w = st.number_input('Enter width of the ROI', min_value=1, value=100) with right_column: y = st.number_input('Enter Y coordinate of top-left corner', min_value=0, value=0) h = st.number_input('Enter height of the ROI', min_value=1, value=100) # Apply adjustments to the frame adjusted_frame = apply_adjustments(frame, grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold, ) # Draw the rectangle on the frame cv2.rectangle(adjusted_frame, (x, y), (x + w, y + h), (255, 0, 0), 2) # Display the adjusted frame st.image(Image.fromarray(adjusted_frame), caption=f'Adjusted Frame at {frame_number}', use_column_width=True) else: st.error(f"Could not read frame number {frame_number}.") st.write('---') # Release the video capture object # vcap.release() return vcap,file_name,total_frames,grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold, x,y,w,h ############################################################################### ############################################################################### ############################################################################### def part_three(vcap,temp_dir,file_name, grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold,x,y,w,h): st.subheader('Part Three: Your object being tracked!') st.write("...please be patient, it may take some time...") vcap.set(cv2.CAP_PROP_POS_FRAMES, 0) # Initialize the object tracker tracker = cv2.TrackerCSRT_create() #initialize positions x_position = [] y_position = [] w_position = [] h_position = [] # Read the first frame success, frame = vcap.read() if success: # Get video properties frame_width = int(vcap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(vcap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = vcap.get(cv2.CAP_PROP_FPS) # Define the codec and create VideoWriter object fourcc = cv2.VideoWriter_fourcc(*'XVID') out = cv2.VideoWriter(f'{temp_dir}/{file_name}_ADJUSTED.avi', fourcc, fps, (frame_width, frame_height)) # Apply adjustments to the first frame adjusted_frame = apply_adjustments(frame, grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold) # Define an initial bounding box bbox = (x, y, w, h) # Initialize tracker with first frame and bounding box tracker.init(adjusted_frame, bbox) tracking_image_box = st.empty() tracking_image_info = st.empty() # Loop over the frames of the video frame_number = 0 while True: # Read a new frame success, frame = vcap.read() if not success: break # Apply adjustments to the frame adjusted_frame = apply_adjustments(frame, grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold) # Update tracker success_tracker, box = tracker.update(adjusted_frame) if success_tracker: # Draw the tracking box (x, y, w, h) = [int(v) for v in box] cv2.rectangle(adjusted_frame, (x, y), (x + w, y + h), (0, 255, 0), 2) # Write the adjusted frame to the output video out.write(adjusted_frame) with tracking_image_info: st.write(f''' Frame number: {frame_number+1} (of {total_frames-1}) (x-position, y-position, width, height) = ''',x,y,w,h) x_position.append(x) y_position.append(y) w_position.append(w) h_position.append(h) # Display the tracked frame with tracking_image_box: st.image(Image.fromarray(adjusted_frame), caption='Tracked Frame', use_column_width=True) frame_number += 1 #Plot x and y over time t = np.linspace(1,len(x_position),len(x_position)) - 1 df = pd.DataFrame({ 'Frame number' : t, 'Horizontal pixel' : x_position, 'Vertical pixel' : y_position, 'Box width' : w_position, 'Box height' : h_position }) else: df=pd.DataFrame() # Release the video capture object vcap.release() out.release() return df ############################################################################### ############################################################################### ############################################################################### def part_four(df,file_name): st.write("---") st.subheader('Part four: RESULTS!') st.write(''' A plot of the horizontal and vertical position of the tracking rectagle (ROI) is displayed in the graph below. You can zoom in and hover over specific parts of the graph as desired. Note, the x-axis is frame number (not time) and the y-axis is pixel position. You will need to convert these to time and distance in your units of interest. ''') st.line_chart( df, x="Frame number", y=["Horizontal pixel", "Vertical pixel"], color=["#FF0000", "#0000FF"] # Optional ) st.write("Please use the two buttons below to download the pixel tracking results as a xlsx file, and a movie of the adjusted video as an avi file.") col_a,col_b = st.columns(2) # Function to convert DataFrame to Excel and return a BytesIO object def to_excel(df): output = BytesIO() with pd.ExcelWriter(output, engine='xlsxwriter') as writer: df.to_excel(writer, index=False, sheet_name='Sheet1') # No need to call writer.save() as it is handled by the context manager processed_data = output.getvalue() return processed_data if len(df)>0: # Download button @st.experimental_fragment def download_excel_results(): st.download_button( label="Download Excel Results File", data=to_excel(df), file_name=f"{file_name}.xlsx", mime="application/vnd.ms-excel" ) with col_a: download_excel_results() with open(f'{temp_dir}/{file_name}_ADJUSTED.avi', 'rb') as file: video_bytes = file.read() @st.experimental_fragment def download_video_results(video_bytes): st.download_button( label="Download Video", data=video_bytes, file_name=f'{file_name}_ADJUSTED.avi', mime="video/avi" ) # video_file.close() # Close the file with col_b: # Function to convert file to a download link if temp_dir: download_video_results(video_bytes) ############################################################################### ############################################################################### ############################################################################### with tempfile.TemporaryDirectory() as temp_dir: # Streamlit app st.title('Cell squashing tracker!') st.write('---') # File uploader uploaded_file = part_one() st.write('---') # If a file is uploaded if uploaded_file is not None: #Image adjustment and ROI vcap,file_name,total_frames,grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold, x,y,w,h = part_two(uploaded_file) # Button to start object tracking if st.button('Start Object Tracking'): #start tracking df = part_three(vcap,temp_dir,file_name, grayscale, contrast, blur, edges, sharpen_amount, subtract_background, lower_threshold, upper_threshold, binarize, binarize_threshold,x,y,w,h) #display/download results part_four(df,file_name) else: st.info("Upload a video file to get started.")