Upload cell_squashing_tracker.py

cell_squashing_tracker.py

@@ -0,0 +1,323 @@
+# -*- 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.")
+