diff --git a/video_cleaner.py b/video_cleaner.py index 5dbc38e..f86d29b 100644 --- a/video_cleaner.py +++ b/video_cleaner.py @@ -83,12 +83,9 @@ class VideoCleaner: prev_gray = gray # --- AI HUMAN DETECTION --- - # We only care about humans if the lights are on (as per your rules) - # OR we can scan anyway if you want to find people in IR mode too. - # Following your rules: "If it's night and no lights, delete". - # So we only need to detect humans to justify keeping the video. - small_frame = cv2.resize(frame, (640, 360)) - results = self.model(small_frame, classes=[0], verbose=False, conf=0.7) + model = self._get_model() + # 直接传原图,省去 CPU resize,开启半精度和指定推理尺寸 + results = model(frame, classes=[0], verbose=False, conf=0.7, half=True, imgsz=640) debugmode=True if len(results[0].boxes) > 0: if is_color: @@ -156,28 +153,34 @@ class VideoCleaner: print(f"Scanning {input_path}...") - # 逐个目录递归处理,避免一次性加载所有文件到内存 + all_videos = [] for root, dirs, files in os.walk(input_path): if "processed" in dirs: dirs.remove("processed") - # 筛选当前目录下的视频文件 - current_dir_videos = [] for file in files: if file.lower().endswith(self.supported_extensions): - current_dir_videos.append(Path(root) / file) - - if not current_dir_videos: - continue + all_videos.append(Path(root) / file) + + if not all_videos: + print("No videos found.") + return - # 对当前目录的文件按时间排序处理 - current_dir_videos.sort(key=lambda x: os.path.getmtime(x)) + # Sort by mtime + all_videos.sort(key=lambda x: os.path.getmtime(x)) + + print(f"Found {len(all_videos)} videos. Starting parallel processing with {self.workers} workers...") + + with ProcessPoolExecutor(max_workers=self.workers) as executor: + # We pass the method and its arguments. + # Note: self will be pickled. Since self.model is None, it should work. + futures = [executor.submit(self.process_video_file, v, processed_dir, input_path) for v in all_videos] - print(f"Processing directory: {root} ({len(current_dir_videos)} files)") - for target_file in current_dir_videos: - # 处理文件。如果 process_video_file 返回 False, - # 通常是因为文件太新(未到 30 天),则跳过当前文件 - self.process_video_file(target_file, processed_dir, input_path) + for future in as_completed(futures): + try: + future.result() + except Exception as e: + print(f"Error in worker process: {e}") if __name__ == "__main__": import argparse @@ -185,6 +188,7 @@ if __name__ == "__main__": parser.add_argument("dir", help="Target directory") parser.add_argument("--days", type=int, default=30) parser.add_argument("--model", type=str, default='yolo26n.pt', help="Path to YOLO model or model name") + parser.add_argument("--workers", type=int, default=4, help="Number of parallel workers") args = parser.parse_args() - cleaner = VideoCleaner(model_path=args.model, age_days=args.days) + cleaner = VideoCleaner(model_path=args.model, age_days=args.days, workers=args.workers) cleaner.scan_and_process(args.dir)