-
Notifications
You must be signed in to change notification settings - Fork 0
/
FrameEnhancer.py
76 lines (61 loc) · 2.33 KB
/
FrameEnhancer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import argparse
import sys
import os
from frame_optimization_methods.opticalFlow import remove_dead_frames as remove_dead_frames_of
from frame_optimization_methods.frameDifference import remove_dead_frames as remove_dead_frames_fd
from frame_optimization_methods.ssim import process_video as process_video_ssim
def main():
parser = argparse.ArgumentParser(
description='Anime Frame Optimizer (AFOptimizer)',
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(
'-of',
'--opticalFlow',
action='store_true',
help='Use the Optical Flow method to remove static frames.')
parser.add_argument(
'-fd',
'--frameDifference',
action='store_true',
help='Use the Frame Difference method to remove static frames.')
parser.add_argument('-ss',
'--ssim',
action='store_true',
help='Use the SSIM method to remove similar frames.')
parser.add_argument('--video',
type=str,
help='Path to the input video file.')
parser.add_argument('--ssim_threshold',
type=float,
default=0.9587,
help='SSIM threshold for frame similarity.')
if len(sys.argv) == 1:
parser.print_help()
sys.exit(1)
args = parser.parse_args()
if not args.video:
print("Error: Please specify the input video file path.")
sys.exit(1)
if not os.path.exists(args.video):
print(f"Error: The video file {args.video} does not exist.")
sys.exit(1)
base_threshold = 10 # Default threshold for frame difference
flow_mag_threshold = 0.4 # Default threshold for optical flow
if args.opticalFlow:
print("Using Optical Flow method...")
remove_dead_frames_of(args.video, flow_mag_threshold)
elif args.frameDifference:
print("Using Frame Difference method...")
remove_dead_frames_fd(args.video, base_threshold)
elif args.ssim:
print("Using SSIM method...")
output_video = os.path.splitext(
args.video)[0] + "_ssim.mp4" # Generating output file name
process_video_ssim(args.video, args.ssim_threshold, output_video)
else:
print(
"Error: No method selected. Use -of for Optical Flow, -fd for Frame Difference, or -ss for SSIM."
)
sys.exit(1)
if __name__ == "__main__":
main()