Video2Crops/video2crops.py

import cv2
import os
import numpy as np
import tkinter as tk
from tkinterdnd2 import DND_FILES, TkinterDnD
from tkinter import filedialog, simpledialog, messagebox
from datetime import datetime
import sys
import webbrowser
import screeninfo

# Constants
v2cversion1 = 0
v2cversion2 = 1
v2cversion3 = 0
v2cversion = f"{v2cversion1}.{v2cversion2}.{v2cversion3}"
frame_step = 100
value_threshold = 11.0
SAVE = True
SHOW = True
processing = False

# Create the root window
root = TkinterDnD.Tk()
root.title("Video2Crops")

# Function to create the output directory and return its path
def create_output_directory(video_filename):
    base_directory = os.getcwd()  # You can change this to specify a different base directory
    video_filename_base, _ = os.path.splitext(os.path.basename(video_filename))
    output_folder_path = os.path.join(base_directory, video_filename_base)

    # Create the output directory if it doesn't exist
    if not os.path.exists(output_folder_path):
        os.makedirs(output_folder_path)

    return output_folder_path

# Function to calculate ROI position based on alignment and size settings
def calculate_roi(video_width, video_height, vertical_alignment, horizontal_alignment, roi_size):
    y1 = (video_height - roi_size) // 2
    x1 = (video_width - roi_size) // 2
    
    if vertical_alignment == "Top":
        y1 = 0
    elif vertical_alignment == "Center":
        y1 = (video_height - roi_size) // 2
    elif vertical_alignment == "Bottom":
        y1 = video_height - roi_size

    if horizontal_alignment == "Left":
        x1 = 0
    elif horizontal_alignment == "Center":
        x1 = (video_width - roi_size) // 2
    elif horizontal_alignment == "Right":
        x1 = video_width - roi_size

    x2 = x1 + roi_size
    y2 = y1 + roi_size

    return x1, y1, x2, y2

# Create a function to process video files
def process_video(video_filenames, main_window):
    global processing
    processing = True

    def write_log(message):
        logfile_path = os.path.join(output_folder, 'logfile.txt')
        with open(logfile_path, 'a') as logfile:
            log_message = f"{datetime.now().strftime('%m-%d_%H:%M:%S')} - {message}\n"
            logfile.write(log_message)
            print(log_message)

    def calc_value(image):
        grayscale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        return np.std(grayscale)

    corner1 = (-1, -1)
    corner2 = (-1, -1)
    drawing = False
    cropped_image = None

    for videofilename in video_filenames:

        main_window.title(f"Video Processing App")

        video_file_src = videofilename

        # Define videofilename_base based on the video file name
        videofilename_base, _ = os.path.splitext(os.path.basename(videofilename))

        cap = cv2.VideoCapture(video_file_src, 0)

        if not cap.isOpened():
            write_log(f"Error: Could not open video file: {video_file_src}")
            sys.exit()

        counter = 0
        frame_cnt = 0
        discarded_frames = 0

        # Create the output folder early on
        output_folder = create_output_directory(videofilename)

        # Create the "Discarded Frames" subfolder within the output folder
        discarded_folder = os.path.join(output_folder, 'Discarded Frames')
        if not os.path.exists(discarded_folder) and SAVE:
            os.makedirs(discarded_folder)

        write_log(f"[DEBUG] Output Folder: {output_folder}\n[DEBUG] Discarded Folder: {discarded_folder}")


        # Initialize ROI parameters from preferences
        roi_size = int(roi_size_var.get())
        vertical_alignment = vertical_alignment_var.get()
        horizontal_alignment = horizontal_alignment_var.get()

        if SHOW:
            vp_title = "Video2Crops Video Player"
            # Get the screen resolution
            screen = screeninfo.get_monitors()[0]
            screen_width, screen_height = screen.width, screen.height

            # Calculate the initial window size as a quarter of the screen resolution
            initial_width = screen_width // 2
            initial_height = screen_height // 2

            # Create the video player window
            cv2.namedWindow(vp_title, cv2.WINDOW_NORMAL)
            cv2.resizeWindow(vp_title, initial_width, initial_height)

        while True:
            ret, frame = cap.read()

            if not ret:
                break

            # Get the video width and height from the frame
            video_width = frame.shape[1]
            video_height = frame.shape[0]

            x1, y1, x2, y2 = calculate_roi(video_width, video_height, vertical_alignment, horizontal_alignment, roi_size)

            cropped_image_status = "Rejected"

            frame_copy = frame.copy()

            if drawing:
                cv2.rectangle(frame_copy, corner1, corner2, (0, 255, 0), 3)

            test = cv2.rectangle(frame.copy(), (x1 - 1, y2 - 1), (x2 + 1, y1 + 1), (255, 0, 0), 4)
            cropped_image = frame[y1:y2, x1:x2]

            if SHOW:
                cv2.imshow(vp_title, test)

            key = cv2.waitKey(10)

            if key == ord('q'):
                break

            if cropped_image is not None and frame_cnt % frame_step == 0:
                if SHOW:
                    ci_title = "Video2Crops Cropped Image"
                    cv2.namedWindow(ci_title, cv2.WINDOW_NORMAL)
                    cv2.imshow(ci_title, cropped_image)

                frame_value = calc_value(cropped_image)
                if frame_value >= value_threshold:
                    cropped_image_status = "Accepted"

                    filename = videofilename_base + '_' + str(counter).zfill(4) + ".png"
                    image_filename = os.path.join(output_folder, filename)

                    if SAVE:
                        cv2.imwrite(image_filename, cropped_image)

                    cropped_image_info = f"Image on frame {frame_cnt} ({cropped_image_status}) - {frame_value} / {value_threshold} - image nr: {str(str(counter)).zfill(4)}"
                    write_log(cropped_image_info)
                    counter += 1
                else:
                    filename = videofilename_base + '_' + str(discarded_frames).zfill(4) + ".png"
                    image_filename = os.path.join(discarded_folder, filename)
                    if SAVE:
                        cv2.imwrite(image_filename, cropped_image)

                    cropped_image_info = f"Image on frame {frame_cnt} ({cropped_image_status}) - {frame_value} / {value_threshold} - image nr: {str(str(counter)).zfill(4)}"
                    write_log(cropped_image_info)
                    discarded_frames += 1

            frame_cnt += 1

        cap.release()
        cv2.destroyAllWindows()

        write_log(f"Video2Crops run complete!\nSaved {counter} frames.\nDiscarded {discarded_frames} frames.")
    
    processing = False


# Create a function to add files to the selected files Listbox
def add_files():
    file_paths = filedialog.askopenfilenames(filetypes=[("Video Files", "*.mp4 *.avi *.mov")])
    if file_paths:
        for file_path in file_paths:
            selected_files_listbox.insert(tk.END, file_path)

# Create a button to remove the selected file from the Listbox
def remove_file():
    selected_index = selected_files_listbox.curselection()
    if selected_index:
        selected_files_listbox.delete(selected_index)

# Create a function to clear the queue
def clear_queue():
    selected_files_listbox.delete(0, tk.END)

# Create a function to cancel processing
def cancel_processing():
    global processing
    if processing:
        processing = False
        messagebox.showinfo("Processing Canceled", "Video processing has been canceled.")
    
# Create a button to process the selected files
def process_files():
    global processing
    if not processing:
        selected_files = selected_files_listbox.get(0, tk.END)
        for file_path in selected_files:
            process_video([file_path], root)

# Create a function to open the preferences window
def open_preferences():

    preferences_window = tk.Toplevel(root)
    preferences_window.title("Preferences")

    def change_frame_step():
        if not processing:
            new_frame_step = simpledialog.askinteger("Frame Step", "Enter the new frame step:")
            if new_frame_step is not None:
                frame_step_var.set(str(new_frame_step))
                global frame_step
                frame_step = new_frame_step

    def change_value_threshold():
        if not processing:
            new_value_threshold = simpledialog.askfloat("Value Threshold", "Enter the new value threshold:")
            if new_value_threshold is not None:
                value_threshold_var.set(str(new_value_threshold))

    def toggle_save():
        if not processing:
            global SAVE
            SAVE = not SAVE

    def toggle_show():
        if not processing:
            global SHOW
            SHOW = not SHOW

    def apply_roi_size():
        if not processing:
            new_roi_size = simpledialog.askinteger("ROI Size", "Enter the new ROI size:")
            if new_roi_size is not None:
                roi_size_var.set(str(new_roi_size))

    # Create preference fields
    frame_step_label = tk.Label(preferences_window, text="Skip n frames:")
    frame_step_label.grid(row=0, column=0, padx=20, pady=10)
    frame_step_value = tk.Label(preferences_window, textvariable=frame_step_var)
    frame_step_value.grid(row=0, column=1, padx=20, pady=10)
    frame_step_button = tk.Button(preferences_window, text="Edit", command=change_frame_step)
    frame_step_button.grid(row=0, column=2, padx=20, pady=10)

    value_threshold_label = tk.Label(preferences_window, text="Discard Threshold:")
    value_threshold_label.grid(row=1, column=0, padx=20, pady=10)
    value_threshold_value = tk.Label(preferences_window, textvariable=value_threshold_var)
    value_threshold_value.grid(row=1, column=1, padx=20, pady=10)
    value_threshold_button = tk.Button(preferences_window, text="Edit", command=change_value_threshold)
    value_threshold_button.grid(row=1, column=2, padx=20, pady=10)

    save_checkbox = tk.Checkbutton(preferences_window, text="Save Output Images", variable=save_var, command=toggle_save)
    save_checkbox.grid(row=2, column=0, columnspan=3, padx=20, pady=10)

    show_checkbox = tk.Checkbutton(preferences_window, text="Display Playback", variable=show_var, command=toggle_show)
    show_checkbox.grid(row=3, column=0, columnspan=3, padx=20, pady=10)

    # Section divider
    divider_label = tk.Label(preferences_window, text="ROI Settings", font=("Helvetica", 12, "bold"))
    divider_label.grid(row=4, column=0, columnspan=3, padx=20, pady=(20, 10))

    # Create ROI fields
    vertical_alignment_label = tk.Label(preferences_window, text="Vertical Alignment:")
    vertical_alignment_label.grid(row=5, column=0, padx=20, pady=10)
    vertical_alignment_var.set("Center")  # Default alignment
    vertical_alignment_menu = tk.OptionMenu(preferences_window, vertical_alignment_var, "Top", "Center", "Bottom")
    vertical_alignment_menu.grid(row=5, column=1, padx=20, pady=10)

    horizontal_alignment_label = tk.Label(preferences_window, text="Horizontal Alignment:")
    horizontal_alignment_label.grid(row=6, column=0, padx=20, pady=10)
    horizontal_alignment_var.set("Center")  # Default alignment
    horizontal_alignment_menu = tk.OptionMenu(preferences_window, horizontal_alignment_var, "Left", "Center", "Right")
    horizontal_alignment_menu.grid(row=6, column=1, padx=20, pady=10)

    roi_size_label = tk.Label(preferences_window, text="ROI Size:")
    roi_size_label.grid(row=7, column=0, padx=20, pady=10)
    roi_size_value = tk.Label(preferences_window, textvariable=roi_size_var)
    roi_size_value.grid(row=7, column=1, padx=20, pady=10)
    roi_size_button = tk.Button(preferences_window, text="Edit", command=apply_roi_size)
    roi_size_button.grid(row=7, column=2, padx=20, pady=10)

# Create a menu
menu = tk.Menu(root)
root.config(menu=menu)

# Variables for preferences
frame_step_var = tk.StringVar()
value_threshold_var = tk.StringVar()
save_var = tk.BooleanVar()
show_var = tk.BooleanVar()

# Variables for ROI
vertical_alignment_var = tk.StringVar()
horizontal_alignment_var = tk.StringVar()
roi_size_var = tk.StringVar()
roi_size_var.set("400")

# Set initial values for preference variables
frame_step_var.set(str(frame_step))
value_threshold_var.set(str(value_threshold))
save_var.set(SAVE)
show_var.set(SHOW)

# Create a "Queue" submenu
queue_menu = tk.Menu(menu, tearoff=0)
menu.add_cascade(label="Queue", menu=queue_menu)
queue_menu.add_command(label="Add file", command=add_files)
queue_menu.add_command(label="Clear queue", command=clear_queue)
queue_menu.add_separator()
queue_menu.add_command(label="Cancel", command=cancel_processing)
queue_menu.add_command(label="Cancel & Clear", command=lambda: [cancel_processing(), clear_queue()])

# Create a "Preferences" submenu
menu.add_command(label="Preferences", command=open_preferences)

# Create an "Info" submenu
info_menu = tk.Menu(menu, tearoff=0)
menu.add_cascade(label="Info", menu=info_menu)
info_menu.add_command(label="Project Folder", command=lambda: webbrowser.open(os.getcwd()))
info_menu.add_separator()
info_menu.add_command(label="About", command=lambda: webbrowser.open("about.md"))
info_menu.add_command(label=f"Version: {str(v2cversion)}")

# Create a label for selected files
label = tk.Label(root, text="Selected Files:")
label.pack(padx=20, pady=10)

# Create a Listbox to display the selected files
selected_files_listbox = tk.Listbox(root, selectmode=tk.SINGLE, exportselection=0)
selected_files_listbox.pack(fill=tk.BOTH, expand=True, padx=20, pady=10)

# Create a button to select and add files to the Listbox
select_files_button = tk.Button(root, text="Select Files", command=add_files)
select_files_button.pack(padx=20, pady=10)

# Create a button to remove the selected file from the Listbox
remove_file_button = tk.Button(root, text="Remove File", command=remove_file)
remove_file_button.pack(padx=20, pady=10)

# Create a button to process the selected files
process_files_button = tk.Button(root, text="Process Files", command=process_files)
process_files_button.pack(padx=20, pady=10)

root.mainloop()