NeuralFalcon commited on
Commit
facc648
·
verified ·
1 Parent(s): 65dc7cb

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +276 -0
app.py ADDED
@@ -0,0 +1,276 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gpu=False
2
+ import easyocr
3
+ reader = easyocr.Reader(['ch_sim','en'],gpu=gpu) # this needs to run only once to load the model into memory
4
+
5
+ import cv2
6
+ import os
7
+ import cv2
8
+ import numpy as np
9
+ import shutil
10
+ from concurrent.futures import ThreadPoolExecutor
11
+ import re
12
+ import subprocess
13
+
14
+ def extract_frames(video_path, output_folder):
15
+ if os.path.exists(output_folder):
16
+ shutil.rmtree(output_folder)
17
+ os.makedirs(output_folder, exist_ok=True)
18
+
19
+ cap = cv2.VideoCapture(video_path)
20
+ frame_count = 0
21
+
22
+ while cap.isOpened():
23
+ ret, frame = cap.read()
24
+ if not ret:
25
+ break # Stop when video ends
26
+
27
+ frame_path = os.path.join(output_folder, f"{frame_count:06d}.png")
28
+ cv2.imwrite(frame_path, frame)
29
+ frame_count += 1
30
+
31
+ cap.release()
32
+ print(f"Extracted {frame_count} frames to {output_folder}")
33
+
34
+
35
+
36
+
37
+
38
+
39
+
40
+ # Initialize text reader
41
+
42
+ def remove_watermark(image, blur_type="strong_gaussian"):
43
+ results = reader.readtext(image) # Detect text regions
44
+
45
+ for (bbox, text, prob) in results:
46
+ top_left = tuple(map(int, bbox[0]))
47
+ bottom_right = tuple(map(int, bbox[2]))
48
+ x1, y1 = top_left
49
+ x2, y2 = bottom_right
50
+ roi = image[y1:y2, x1:x2]
51
+
52
+ if blur_type == "strong_gaussian":
53
+ blurred_roi = cv2.GaussianBlur(roi, (25, 25), 50)
54
+ elif blur_type == "pixelation":
55
+ h, w = roi.shape[:2]
56
+ temp = cv2.resize(roi, (8, 8), interpolation=cv2.INTER_LINEAR)
57
+ blurred_roi = cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)
58
+ elif blur_type == "median":
59
+ blurred_roi = cv2.medianBlur(roi, 21)
60
+ elif blur_type == "motion":
61
+ size = 25
62
+ kernel = np.zeros((size, size))
63
+ kernel[:, size//2] = 1
64
+ kernel = kernel / kernel.sum()
65
+ blurred_roi = cv2.filter2D(roi, -1, kernel)
66
+ elif blur_type == "bilateral":
67
+ blurred_roi = cv2.bilateralFilter(roi, d=15, sigmaColor=75, sigmaSpace=75)
68
+ elif blur_type == "box":
69
+ blurred_roi = cv2.blur(roi, (25, 25))
70
+ elif blur_type == "stacked":
71
+ temp = cv2.GaussianBlur(roi, (15, 15), 25)
72
+ blurred_roi = cv2.medianBlur(temp, 15)
73
+ elif blur_type == "adaptive":
74
+ gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
75
+ _, mask = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY)
76
+ blurred = cv2.GaussianBlur(roi, (25, 25), 25)
77
+ blurred_roi = np.where(mask[..., None] > 0, blurred, roi)
78
+ else:
79
+ blurred_roi = cv2.GaussianBlur(roi, (25, 25), 50)
80
+
81
+ image[y1:y2, x1:x2] = blurred_roi
82
+ return image
83
+ # return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
84
+
85
+
86
+
87
+
88
+ def process_frame(frame_path, save_path):
89
+ image = cv2.imread(frame_path)
90
+
91
+ if image is None:
92
+ print(f"Failed to load: {frame_path}") # Debugging step
93
+ return
94
+
95
+ no_watermark_image = remove_watermark(image, blur_type="median")
96
+
97
+ output_file = os.path.join(save_path, os.path.basename(frame_path))
98
+ success = cv2.imwrite(output_file, no_watermark_image)
99
+
100
+ if not success:
101
+ print(f"Failed to save: {output_file}") # Debugging step
102
+
103
+ def batch_process(batch_size=100):
104
+ input_folder = "./frames"
105
+ output_folder = "./clean"
106
+
107
+ if os.path.exists(output_folder):
108
+ shutil.rmtree(output_folder)
109
+ os.makedirs(output_folder, exist_ok=True)
110
+
111
+ frame_paths = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.endswith((".jpg", ".png"))]
112
+
113
+ with ThreadPoolExecutor() as executor:
114
+ executor.map(process_frame, frame_paths, [output_folder] * len(frame_paths))
115
+
116
+ print(f"Processing complete! {len(frame_paths)} frames saved to {output_folder}")
117
+
118
+
119
+
120
+
121
+ def get_video_fps(video_path):
122
+ """Extract FPS from the original video."""
123
+ cap = cv2.VideoCapture(video_path)
124
+ if not cap.isOpened():
125
+ return None
126
+ fps = cap.get(cv2.CAP_PROP_FPS)
127
+ cap.release()
128
+ return fps
129
+
130
+ def sorted_files(directory):
131
+ """Returns a list of sorted .png files based on numeric order."""
132
+ files = [f for f in os.listdir(directory) if f.endswith(".png")]
133
+ files.sort(key=lambda f: int(re.search(r'\d+', f).group()) if re.search(r'\d+', f) else float('inf'))
134
+ return [os.path.join(directory, f) for f in files]
135
+
136
+ def create_video_chunks(frame_dir, output_dir, fps, batch_size=100):
137
+ """Creates chunked videos from frames in batches."""
138
+
139
+ # Remove old "chunks" folder if exists
140
+ if os.path.exists(output_dir):
141
+ shutil.rmtree(output_dir)
142
+ os.makedirs(output_dir, exist_ok=True)
143
+
144
+ sorted_images = sorted_files(frame_dir)
145
+
146
+ total_chunks = (len(sorted_images) // batch_size) + (1 if len(sorted_images) % batch_size else 0)
147
+
148
+ for i in range(total_chunks):
149
+ chunk_frames = sorted_images[i * batch_size:(i + 1) * batch_size]
150
+ if not chunk_frames:
151
+ continue
152
+
153
+ chunk_folder = os.path.join(output_dir, f"chunk_{i+1}")
154
+ os.makedirs(chunk_folder, exist_ok=True)
155
+
156
+ # Copy frames to a temp folder
157
+ for j, frame in enumerate(chunk_frames):
158
+ frame_dest = os.path.join(chunk_folder, f"{j:05d}.png") # Zero-padded filenames
159
+ shutil.copy(frame, frame_dest)
160
+
161
+ # Generate video from frames
162
+ chunk_output = os.path.join(output_dir, f"{i+1}.mp4")
163
+ ffmpeg_cmd = f'ffmpeg -y -framerate {fps} -i "{chunk_folder}/%05d.png" -c:v libx264 -pix_fmt yuv420p "{chunk_output}"'
164
+ subprocess.run(ffmpeg_cmd, shell=True, check=True)
165
+
166
+ # Cleanup temp chunk folder
167
+ shutil.rmtree(chunk_folder)
168
+
169
+ print(f"✅ All {total_chunks} video chunks created in {output_dir}")
170
+
171
+ def vido_chunks(video_path):
172
+ # Extract original FPS
173
+ fps = get_video_fps(video_path)
174
+ if fps is None:
175
+ raise ValueError("Failed to retrieve FPS from video.")
176
+
177
+ # Define folders
178
+ frame_dir = "./clean"
179
+ output_dir = "./chunks"
180
+ # Process frames into video chunks
181
+ create_video_chunks(frame_dir, output_dir, fps, batch_size=100)
182
+
183
+ import os
184
+ import re
185
+ import uuid
186
+
187
+ def sanitize_file(file_path):
188
+ folder = os.path.dirname(file_path)
189
+ text, ext = os.path.splitext(os.path.basename(file_path))
190
+
191
+ # Keep alphabets, spaces, and underscores only
192
+ text = re.sub(r'[^a-zA-Z_ ]', '', text)
193
+ text = text.lower().strip()
194
+ text = text.replace(" ", "_")
195
+
196
+ # Truncate or handle empty text
197
+ truncated_text = text[:20] if len(text) > 20 else text if len(text) > 0 else "empty"
198
+
199
+ # Generate a random string for uniqueness
200
+ random_string = uuid.uuid4().hex[:8].upper()
201
+
202
+ # Construct the new file name
203
+ # file_name = f"{folder}/{truncated_text}_{random_string}{ext}"
204
+ file_name = f"{truncated_text}_{random_string}{ext}"
205
+ return file_name
206
+ def upload_file(video_path):
207
+ os.makedirs("./upload",exist_ok=True)
208
+ new_path=sanitize_file(video_path)
209
+ new_path=f"./upload/{new_path}"
210
+ shutil.copy(video_path,new_path)
211
+ return new_path
212
+
213
+
214
+
215
+ import os
216
+ import re
217
+ import subprocess
218
+ def sorted_video_files(directory):
219
+ """Returns a list of full paths of .mp4 files sorted by the numeric part of the filename."""
220
+ files = [f for f in os.listdir(directory) if f.endswith(".mp4")]
221
+
222
+ # Extract the numeric part using regex and sort
223
+ files.sort(key=lambda f: int(re.search(r'\d+', f).group()) if re.search(r'\d+', f) else float('inf'))
224
+
225
+ # Convert filenames to full paths
226
+ full_paths = [os.path.join(directory, f) for f in files]
227
+
228
+ return full_paths
229
+
230
+ def marge_video(gpu=True):
231
+ os.makedirs("./result/",exist_ok=True)
232
+ output_path=f"./result/no_water_mark.mp4"
233
+ video_list=sorted_video_files("./chunks")
234
+ with open("./join.txt", "w") as f:
235
+ for video in video_list:
236
+ f.write(f"file '{video}'\n")
237
+ if gpu:
238
+ join_command = f'ffmpeg -hwaccel cuda -f concat -safe 0 -i ./join.txt -c copy "{output_path}" -y'
239
+ else:
240
+ join_command = f'ffmpeg -f concat -safe 0 -i ./join.txt -c copy "{output_path}" -y'
241
+ subprocess.run(join_command, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
242
+ return output_path
243
+ def recover_audio(upload_path):
244
+ output_path=f"./result/no_water_mark.mp4"
245
+ audio_path="./upload/temp.wav"
246
+ save_path=upload_path.replace(".mp4","_no_watermark.mp4")
247
+ var=os.system(f"ffmpeg -i {upload_path} -q:a 0 -map a {audio_path} -y")
248
+ if var==0:
249
+ var2=os.system(f"ffmpeg -i {output_path} -i {audio_path} -c:v copy -map 0:v:0 -map 1:a:0 -shortest {save_path} -y")
250
+ if var2==0:
251
+ return save_path
252
+ return None
253
+ def video_watermark_remover(video_path):
254
+ global gpu
255
+ upload_path=upload_file(video_path)
256
+ extract_frames(upload_path, "./frames")
257
+ video_path = "/content/face.mp4"
258
+ vido_chunks(upload_path)
259
+ marge_video(gpu=gpu)
260
+ save_path=recover_audio(upload_path)
261
+ return save_path
262
+
263
+
264
+ import gradio as gr
265
+ def gradio_interface(video_file):
266
+ return video_watermark_remover(video_file)
267
+
268
+ demo = gr.Interface(
269
+ fn=gradio_interface,
270
+ inputs=gr.Video(label="Upload Video"),
271
+ outputs=gr.File(label="Processed Video"),
272
+ title="Video Watermark Remover",
273
+ description="Upload a video, and this tool will remove watermarks using blurring techniques."
274
+ )
275
+
276
+ demo.launch()