Spaces:
Running
Running
| import cv2 | |
| from ultralytics import YOLO | |
| from moviepy.editor import VideoFileClip | |
| model = YOLO("yolov8-face-hf.pt") | |
| # ------------------ IMAGE FUNCTION ------------------ # | |
| def blur_faces_image(image): | |
| results = model(image) | |
| for result in results: | |
| for box in result.boxes.xyxy.cpu().numpy(): | |
| x1, y1, x2, y2 = map(int, box) | |
| x1, y1 = max(0, x1), max(0, y1) | |
| x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) | |
| face = image[y1:y2, x1:x2] | |
| if face.size > 0: | |
| blurred = cv2.GaussianBlur(face, (101, 101), 30) | |
| image[y1:y2, x1:x2] = blurred | |
| return image | |
| # ------------------ VIDEO FUNCTION ------------------ # | |
| def blur_faces_video(input_path, output_path, update_callback=None): | |
| cap = cv2.VideoCapture(input_path) | |
| if not cap.isOpened(): | |
| raise IOError("Cannot open video") | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| width = width if width % 2 == 0 else width - 1 | |
| height = height if height % 2 == 0 else height - 1 | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| temp_video_path = "blurred_temp.mp4" | |
| out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height)) | |
| frame_num = 0 | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame = cv2.resize(frame, (width, height)) | |
| results = model(frame, verbose=False)[0] | |
| h, w = frame.shape[:2] | |
| for box in results.boxes.xyxy: | |
| x1, y1, x2, y2 = map(int, box) | |
| x1 = max(0, min(w, x1)) | |
| x2 = max(0, min(w, x2)) | |
| y1 = max(0, min(h, y1)) | |
| y2 = max(0, min(h, y2)) | |
| face = frame[y1:y2, x1:x2] | |
| if face.size > 0: | |
| blur_size = max(55, ((x2 - x1) // 2) | 1) | |
| if blur_size % 2 == 0: | |
| blur_size += 1 | |
| blurred = cv2.GaussianBlur(face, (blur_size, blur_size), 0) | |
| frame[y1:y2, x1:x2] = blurred | |
| out.write(frame) | |
| frame_num += 1 | |
| if update_callback and total_frames > 0: | |
| update_callback(frame_num / total_frames) | |
| cap.release() | |
| out.release() | |
| # Merge audio | |
| try: | |
| if update_callback: | |
| update_callback(0.99) | |
| original = VideoFileClip(input_path) | |
| processed = VideoFileClip(temp_video_path).set_audio(original.audio) | |
| processed.write_videofile( | |
| output_path, | |
| codec="libx264", | |
| audio_codec="aac", | |
| threads=1, | |
| logger=None | |
| ) | |
| if update_callback: | |
| update_callback(1.0) | |
| return output_path | |
| except Exception as e: | |
| print("Audio merging failed:", e) | |
| return temp_video_path | |