Upload open source code of MTFL model

28e129b verified 12 months ago

5.89 kB

"""
Annotation Parser Script
This script is used for annotating video files with anomaly information, such as the number of frames containing anomalies
and the start-end frame couples where anomalies occur. It provides a graphical interface for replaying videos, pausing,
marking frames with anomalies, and undoing markings.
Usage:
    python video_annotator.py --root_dir <root_directory> --video_subdir <video_subdirectory> --annotation_file <output_annotation_file> --anomaly_type <anomaly_class_name>
Parameters:
    --root_dir: Path to the root directory containing the video files.
    --video_subdir: Subdirectory within the root directory where video files are located.
    --annotation_file: Path to the output annotation file where annotated information will be saved.
    --anomaly_type: The class name or type of anomalies to be annotated.
Note:
    Please run this in your local environment and then transfer the annotation file to the server for
    further tasks. The server environment may have plugin issues due to the opencv-python.
"""
import cv2
import sys
import os
import argparse


def get_args():
    parser = argparse.ArgumentParser(description="Annotation Parser")
    # io
    parser.add_argument('--root_dir', type=str, default="H:\\Projects\\VAD\\Test",
                        help="path to root directory")
    parser.add_argument('--video_subdir', type=str, default="Anomaly_videos",
                        help="path to video subfolder")
    parser.add_argument('--annotation_file', type=str, default="H:\\Projects\\VAD\\Test\\Anomaly_videos.txt",
                        help="the output annotation file")
    parser.add_argument('--anomaly_type', type=str, default='Abnormal', help="The class name of videos.")

    return parser.parse_args()


def get_files(folder_path):
    # Store a list of video files
    video_files = []

    # Traverse the folder
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith(('.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv')):
                video_files.append(os.path.join(root, file))

    return video_files


if __name__ == "__main__":
    args = get_args()
    # Loop over all files in directory
    video_dir = os.path.join(args.root_dir, args.video_subdir)
    files = get_files(video_dir)

    for filename in files:
        if not filename.endswith(".mp4"):
            continue

        cap = cv2.VideoCapture(filename)
        num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print(f"{filename}: {num_frames}")

        if not cap.isOpened():
            print("Error opening video")
            sys.exit(1)

        mem = []
        actions = []

        counter = 0
        max_frame = 0
        marked_frame = 0

        playing = True
        marking = False

        while cap.isOpened():
            if playing:
                if counter == max_frame:
                    ret, frame = cap.read()
                    if ret:
                        if len(mem) == 200:
                            mem.pop(0)
                        mem.append(frame)
                    else:
                        break
                    counter += 1
                    max_frame += 1

                else:
                    frame = mem[counter - max_frame - 1]
                    counter += 1
            else:
                frame = mem[counter - max_frame - 1]
                cv2.putText(frame, f"{len(mem) + counter - max_frame}", (50, 50), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                            fontScale=1, color=(255, 255, 255))

            if marking:
                cv2.circle(frame, (int(frame.shape[1]/2), 50), 10, (0, 0, 255), -1)

            cv2.imshow("Replay", frame)
            key = cv2.waitKey(60)
            if key & 0xFF == ord(' '):  # Press space to pause
                playing = not playing
            elif key & 0xFF == ord('a') and not playing:  # a to reverse 10 frames
                if (max_frame - counter) < len(mem) - 11:
                    counter -= 10
            elif key & 0xFF == ord('d') and not playing:  # d to skip 10 frames
                if (max_frame - counter) > 10:
                    counter += 10
            elif key & 0xFF == ord(',') and not playing:  # , to reverse 1 frame
                if (max_frame - counter) < len(mem) - 1:
                    counter -= 1
            elif key & 0xFF == ord('.') and not playing:  # . to skip 1 frame
                if (max_frame - counter) > 0:
                    counter += 1
            elif key & 0xFF == ord('m'):  # m to mark this frame
                if marking:
                    actions.append((marked_frame, counter))
                else:
                    marked_frame = counter
                marking = not marking
                print(f'marking {counter}, waiting for the end frame: {marking}')
            elif key & 0xFF == ord('z'):  # z to undo last marking
                print(f'cancel last marking')
                if not marking:
                    actions.pop(-1)
                marking = not marking
            elif key & 0xFF == ord('s'):  # s to skip this video
                break
            elif key & 0xFF == ord('q'):  # q to quit
                sys.exit(1)

        # the annotation format:
        # relative path to video dir    anomaly type      num_frames     start-end-couples
        num_action = len(actions)
        dir_and_file = os.path.split(filename)
        relative_path = os.path.relpath(filename, args.root_dir)
        line = f"{relative_path} {args.anomaly_type} {num_frames}"
        for i in range(num_action):
            line += f" {actions[i][0]} {actions[i][1]}"
        line += '\n'

        with open(args.annotation_file, "a") as f:
            f.write(line)

        cap.release()
        cv2.destroyWindow("Replay")

    cv2.destroyAllWindows()

	"""
	Annotation Parser Script

	This script is used for annotating video files with anomaly information, such as the number of frames containing anomalies
	and the start-end frame couples where anomalies occur. It provides a graphical interface for replaying videos, pausing,
	marking frames with anomalies, and undoing markings.

	Usage:
	python video_annotator.py --root_dir <root_directory> --video_subdir <video_subdirectory> --annotation_file <output_annotation_file> --anomaly_type <anomaly_class_name>

	Parameters:
	--root_dir: Path to the root directory containing the video files.
	--video_subdir: Subdirectory within the root directory where video files are located.
	--annotation_file: Path to the output annotation file where annotated information will be saved.
	--anomaly_type: The class name or type of anomalies to be annotated.

	Note:
	Please run this in your local environment and then transfer the annotation file to the server for
	further tasks. The server environment may have plugin issues due to the opencv-python.
	"""
	import cv2
	import sys
	import os
	import argparse


	def get_args():
	parser = argparse.ArgumentParser(description="Annotation Parser")
	# io
	parser.add_argument('--root_dir', type=str, default="H:\\Projects\\VAD\\Test",
	help="path to root directory")
	parser.add_argument('--video_subdir', type=str, default="Anomaly_videos",
	help="path to video subfolder")
	parser.add_argument('--annotation_file', type=str, default="H:\\Projects\\VAD\\Test\\Anomaly_videos.txt",
	help="the output annotation file")
	parser.add_argument('--anomaly_type', type=str, default='Abnormal', help="The class name of videos.")

	return parser.parse_args()


	def get_files(folder_path):
	# Store a list of video files
	video_files = []

	# Traverse the folder
	for root, dirs, files in os.walk(folder_path):
	for file in files:
	if file.lower().endswith(('.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv')):
	video_files.append(os.path.join(root, file))

	return video_files


	if __name__ == "__main__":
	args = get_args()
	# Loop over all files in directory
	video_dir = os.path.join(args.root_dir, args.video_subdir)
	files = get_files(video_dir)

	for filename in files:
	if not filename.endswith(".mp4"):
	continue

	cap = cv2.VideoCapture(filename)
	num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
	print(f"{filename}: {num_frames}")

	if not cap.isOpened():
	print("Error opening video")
	sys.exit(1)

	mem = []
	actions = []

	counter = 0
	max_frame = 0
	marked_frame = 0

	playing = True
	marking = False

	while cap.isOpened():
	if playing:
	if counter == max_frame:
	ret, frame = cap.read()
	if ret:
	if len(mem) == 200:
	mem.pop(0)
	mem.append(frame)
	else:
	break
	counter += 1
	max_frame += 1

	else:
	frame = mem[counter - max_frame - 1]
	counter += 1
	else:
	frame = mem[counter - max_frame - 1]
	cv2.putText(frame, f"{len(mem) + counter - max_frame}", (50, 50), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
	fontScale=1, color=(255, 255, 255))

	if marking:
	cv2.circle(frame, (int(frame.shape[1]/2), 50), 10, (0, 0, 255), -1)

	cv2.imshow("Replay", frame)
	key = cv2.waitKey(60)
	if key & 0xFF == ord(' '): # Press space to pause
	playing = not playing
	elif key & 0xFF == ord('a') and not playing: # a to reverse 10 frames
	if (max_frame - counter) < len(mem) - 11:
	counter -= 10
	elif key & 0xFF == ord('d') and not playing: # d to skip 10 frames
	if (max_frame - counter) > 10:
	counter += 10
	elif key & 0xFF == ord(',') and not playing: # , to reverse 1 frame
	if (max_frame - counter) < len(mem) - 1:
	counter -= 1
	elif key & 0xFF == ord('.') and not playing: # . to skip 1 frame
	if (max_frame - counter) > 0:
	counter += 1
	elif key & 0xFF == ord('m'): # m to mark this frame
	if marking:
	actions.append((marked_frame, counter))
	else:
	marked_frame = counter
	marking = not marking
	print(f'marking {counter}, waiting for the end frame: {marking}')
	elif key & 0xFF == ord('z'): # z to undo last marking
	print(f'cancel last marking')
	if not marking:
	actions.pop(-1)
	marking = not marking
	elif key & 0xFF == ord('s'): # s to skip this video
	break
	elif key & 0xFF == ord('q'): # q to quit
	sys.exit(1)

	# the annotation format:
	# relative path to video dir anomaly type num_frames start-end-couples
	num_action = len(actions)
	dir_and_file = os.path.split(filename)
	relative_path = os.path.relpath(filename, args.root_dir)
	line = f"{relative_path} {args.anomaly_type} {num_frames}"
	for i in range(num_action):
	line += f" {actions[i][0]} {actions[i][1]}"
	line += '\n'

	with open(args.annotation_file, "a") as f:
	f.write(line)

	cap.release()
	cv2.destroyWindow("Replay")

	cv2.destroyAllWindows()