ytcollector 1.0.8__py3-none-any.whl → 1.1.1__py3-none-any.whl

ytcollector/cli.py

@@ -1,234 +1,202 @@
 #!/usr/bin/env python3
-"""
-SBS Dataset Collector CLI (Updated)
-"""
-import argparse
-import logging
-from pathlib import Path
+"""YouTube 콘텐츠 수집기 - CLI 모듈"""
 
-# Package import modified to 'downloader'
-from .config import TASK_CLASSES, VALID_TASKS, get_paths
-from .utils import ensure_dir, get_url_file_path
+import argparse
+import os
+import sys
 
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+from .config import CATEGORY_NAMES
+from .downloader import YouTubeDownloader
+from .analyzer import check_dependencies
 
 
-def init_project(args):
-    """프로젝트 디렉토리 초기화"""
-    base_dir = Path(args.dir) if args.dir else Path.cwd()
-    
-    # 1. 태스크별 폴더 및 youtube_url.txt 생성
-    for task in VALID_TASKS:
-        # get_url_file_path 내부에서 ensure_dir 호출로 폴더 생성됨
-        txt_path = get_url_file_path(base_dir, task)
-        
-        if not txt_path.exists():
-            txt_path.write_text(
-                "task_type,url,timestamp_min,timestamp_sec,description\n"
-                f"{task},https://www.youtube.com/watch?v=EXAMPLE,2,30,샘플\n",
-                encoding='utf-8'
-            )
-            
-    # 2. config.py의 get_paths 로직에 따른 폴더들 생성 확인
-    paths = get_paths(base_dir)
-    ensure_dir(paths['outputs']) 
-    
-    print(f"✓ Project initialized at: {base_dir}")
-    print(f"  - Add URLs to: urls/<task>/youtube_url.txt")
-    print(f"  - Videos will be saved to configured OUTPUT_DIR (or video/ folder)")
+def create_parser():
+    """CLI 인자 파서 생성"""
+    parser = argparse.ArgumentParser(
+        prog='ytcollector',
+        description='YouTube 콘텐츠 수집기 - 얼굴, 번호판, 타투, 텍스트 감지',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+예시:
+  ytcollector -c face                    # 얼굴 카테고리 5개
+  ytcollector -c face text -n 10         # 얼굴, 텍스트 각 10개
+  ytcollector -c face --fast             # 고속 모드 (병렬 다운로드)
+  ytcollector -c face --fast -w 5        # 5개 동시 다운로드
+  ytcollector -c license_plate -d 5      # 번호판, 최대 5분
+  
+  # 짧은 명령어도 사용 가능
+  ytc -c face -n 3
+        """
+    )
 
+    parser.add_argument(
+        '-c', '--categories',
+        nargs='+',
+        choices=['face', 'license_plate', 'tattoo', 'text'],
+        default=['face'],
+        help='수집할 카테고리 (기본: face)'
+    )
+    parser.add_argument(
+        '-n', '--count',
+        type=int,
+        default=5,
+        help='카테고리당 다운로드 수 (기본: 5)'
+    )
+    parser.add_argument(
+        '-d', '--duration',
+        type=int,
+        default=3,
+        help='최대 영상 길이(분) (기본: 3)'
+    )
+    parser.add_argument(
+        '-o', '--output',
+        type=str,
+        default=".",
+        help='저장 경로 (기본: 현재 폴더)'
+    )
+    parser.add_argument(
+        '--fast',
+        action='store_true',
+        help='고속 모드 (병렬 다운로드, 딜레이 최소화)'
+    )
+    parser.add_argument(
+        '-w', '--workers',
+        type=int,
+        default=3,
+        help='병렬 다운로드 수 (기본: 3, --fast 필요)'
+    )
+    parser.add_argument(
+        '--proxy',
+        type=str,
+        default=None,
+        help='프록시 (예: http://proxy:8080)'
+    )
+    parser.add_argument(
+        '-v', '--version',
+        action='version',
+        version='%(prog)s 1.1.1'
+    )
+    parser.add_argument(
+        '--check-deps',
+        action='store_true',
+        help='의존성 확인 후 종료'
+    )
 
-def run_download(args):
-    """TXT 파일에서 영상 다운로드"""
-    from .downloader import download_from_txt  # Changed function name
-    
-    base_dir = Path(args.dir) if args.dir else Path.cwd()
-    
-    # argparse에서 nargs='+'로 받아오면 args.task는 항상 리스트
-    tasks = args.task if isinstance(args.task, list) else [args.task]
-    
-    total_success = 0
-    total_processed = 0
-    
-    for task in tasks:
-        logger.info(f"=== Processing Task: {task} ===")
-        
-        # 파일 경로: video/{task}/youtube_url.txt
-        txt_file = get_url_file_path(base_dir, task)
-        
-        if not txt_file.exists():
-            logger.error(f"URL file not found: {txt_file}")
-            logger.info(f"Skipping {task}. Run 'ytcollector init' first.")
-            continue
-        
-        logger.info(f"Starting{' fast' if args.fast else ''} download for task: {task}")
-        
-        if args.fast:
-            from .downloader import download_from_txt_parallel
-            results = download_from_txt_parallel(txt_file, task, base_dir, max_count=args.count, skip_verify=args.skip_verify)
-        else:
-            results = download_from_txt(txt_file, task, base_dir, max_count=args.count, skip_verify=args.skip_verify)
-        
-        success_count = sum(1 for r in results if r.get('success'))
-        total_success += success_count
-        total_processed += len(results)
-        
-        print(f"✓ Task '{task}' complete: {success_count}/{len(results)} successful")
+    return parser
+
+
+def run(
+    categories=None,
+    count=5,
+    duration=3,
+    output=None,
+    fast=False,
+    workers=3,
+    proxy=None,
+    quiet=False
+):
+    """
+    프로그래밍 방식으로 수집기 실행
+    
+    Args:
+        categories: 카테고리 리스트 (예: ['face', 'text'])
+        count: 카테고리당 다운로드 수
+        duration: 최대 영상 길이(분)
+        output: 저장 경로
+        fast: 고속 모드 여부
+        workers: 병렬 다운로드 수
+        proxy: 프록시 URL
+        quiet: 조용한 모드 (출력 최소화)
         
-    print(f"\n✓ All tasks complete: {total_success}/{total_processed} successful total")
-
-
-def run_download_single(args):
-    """단일 URL 다운로드"""
-    from .downloader import VideoDownloader
-    
-    base_dir = Path(args.dir) if args.dir else Path.cwd()
-    downloader = VideoDownloader(args.task, base_dir)
-    
-    try:
-        output_path, metadata = downloader.download_clip_at_timestamp(
-            url=args.url,
-            timestamp_min=args.timestamp_min,
-            timestamp_sec=args.timestamp_sec
-        )
-        status = "Cached" if metadata.get('cached') else "Downloaded"
-        print(f"✓ {status}: {output_path}")
-        if not metadata.get('cached'):
-            print(f"  Clip duration: {metadata['clip_duration']}s")
-            
-    except Exception as e:
-        logger.error(f"Download failed: {e}")
-
-
-def run_verify(args):
-    """클립 영상 검증"""
-    from .verifier import verify_clip, batch_verify
+    Returns:
+        dict: 카테고리별 성공 다운로드 수
+    """
+    if categories is None:
+        categories = ['face']
+    
+    if output is None:
+        output = "."
+    
+    # 의존성 체크
+    missing = check_dependencies()
+    if missing and not quiet:
+        print(f"⚠ 분석 기능을 위해 설치 필요: pip install {' '.join(missing)}")
+    
+    # 다운로더 생성
+    downloader = YouTubeDownloader(
+        output_path=output,
+        max_duration=duration * 60,
+        proxy=proxy,
+        fast_mode=fast,
+        workers=workers
+    )
     
-    base_dir = Path(args.dir) if args.dir else Path.cwd()
+    results = {}
+    for category in categories:
+        count_success = downloader.collect(category, count)
+        results[category] = count_success
     
-    if args.video:
-        video_path = Path(args.video)
-        result = verify_clip(video_path, args.task, base_dir)
-        print_verification_result(result)
-    else:
-        # 폴더 경로: video/{task}/
-        clips_dir = base_dir / "video" / args.task
-        if not clips_dir.exists():
-            logger.error(f"Video directory not found: {clips_dir}")
-            return
-        
-        results = batch_verify(clips_dir, args.task, base_dir)
-        valid_count = sum(1 for r in results if r.get('is_valid'))
-        print(f"✓ Verification complete: {valid_count}/{len(results)} valid")
+    return results
 
 
-def run_pipeline(args):
-    """다운로드 + 검증 전체 파이프라인"""
-    print(f"=== Starting pipeline for task: {args.task} ===")
-    
-    run_download(args)
-    
-    if args.verify:
-        print("\n--- Running verification ---")
-        run_verify(args)
-    
-    print("=== Pipeline complete ===")
-
+def main(args=None):
+    """CLI 메인 엔트리포인트"""
+    parser = create_parser()
+    parsed_args = parser.parse_args(args)
 
-def print_verification_result(result: dict):
-    """검증 결과 출력"""
-    summary = result.get('summary', {})
-    
-    print("\n" + "="*50)
-    print(f"Video: {Path(result['video_path']).name}")
-    print(f"Task: {result['task_type']}")
-    print(f"Classes: {result['classes']}")
-    print("-"*50)
-    print(f"Duration: {summary.get('duration_sec', 0):.1f}s")
-    print(f"Frames with detection: {summary.get('frames_with_detection', 0)}")
-    print(f"Detection rate: {summary.get('detection_rate', 0):.1%}")
-    print(f"Valid: {'✓ YES' if result.get('is_valid') else '✗ NO'}")
-    print("="*50)
-
-
-def list_tasks(args):
-    """태스크 목록 출력"""
-    print("\nAvailable Tasks and YOLO-World Classes:")
-    print("-" * 50)
-    for task, classes in TASK_CLASSES.items():
-        print(f"\n{task}:")
-        for cls in classes:
-            print(f"  - {cls}")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description='Downloader - SBS Dataset Collector',
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  downloader init                           # 프로젝트 초기화
-  downloader download --task face           # 텍스트 파일에서 다운로드
-  downloader verify --task face             # YOLO 검증
-        """
+    # 의존성 확인 모드
+    if parsed_args.check_deps:
+        missing = check_dependencies()
+        if missing:
+            print("⚠ 누락된 의존성:")
+            for dep in missing:
+                print(f"  - {dep}")
+            print(f"\n설치: pip install {' '.join(missing)}")
+            sys.exit(1)
+        else:
+            print("✅ 모든 의존성이 설치되어 있습니다.")
+            sys.exit(0)
+
+    # 시작 메시지
+    print("\n" + "=" * 60)
+    print("YouTube 콘텐츠 수집기")
+    print("=" * 60)
+    print(f"카테고리: {', '.join([CATEGORY_NAMES[c] for c in parsed_args.categories])}")
+    print(f"개수: 카테고리당 {parsed_args.count}개")
+    print(f"최대길이: {parsed_args.duration}분")
+    print(f"저장경로: {parsed_args.output}")
+    if parsed_args.fast:
+        print(f"모드: ⚡ 고속 (병렬 {parsed_args.workers}개)")
+    if parsed_args.proxy:
+        print(f"프록시: {parsed_args.proxy}")
+
+    # 의존성 체크
+    missing = check_dependencies()
+    if missing:
+        print(f"\n⚠ 분석 기능을 위해 설치 필요:")
+        print(f"  pip install {' '.join(missing)}")
+
+    # 수집 실행
+    results = run(
+        categories=parsed_args.categories,
+        count=parsed_args.count,
+        duration=parsed_args.duration,
+        output=parsed_args.output,
+        fast=parsed_args.fast,
+        workers=parsed_args.workers,
+        proxy=parsed_args.proxy
     )
-    
-    parser.add_argument('--dir', '-d', help='Project directory (default: current)')
-    
-    subparsers = parser.add_subparsers(dest='command', help='Commands')
-    
-    # Init
-    init_parser = subparsers.add_parser('init', help='Initialize project directory')
-    
-    # Download
-    download_parser = subparsers.add_parser('download', help='Download from youtube_url.txt')
-    download_parser.add_argument('--task', '-t', required=True, nargs='+', choices=VALID_TASKS, help='One or more tasks (e.g. face tattoo)')
-    download_parser.add_argument('--count', '-n', type=int, help='Max videos to collect (default: 1000)')
-    download_parser.add_argument('--fast', action='store_true', help='Enable fast parallel downloading')
-    download_parser.add_argument('--skip-verify', '-S', action='store_true', help='Skip YOLO verification and save all clips')
-    
-    # Download single
-    single_parser = subparsers.add_parser('download-single', help='Download single video')
-    single_parser.add_argument('--task', '-t', required=True, choices=VALID_TASKS)
-    single_parser.add_argument('--url', '-u', required=True, help='YouTube URL')
-    single_parser.add_argument('--timestamp-min', '-m', type=int, required=True)
-    single_parser.add_argument('--timestamp-sec', '-s', type=int, required=True)
-    
-    # Verify
-    verify_parser = subparsers.add_parser('verify', help='Verify with YOLO-World')
-    verify_parser.add_argument('--task', '-t', required=True, nargs='+', choices=VALID_TASKS)
-    verify_parser.add_argument('--video', '-v', help='Specific video file')
-    
-    # Pipeline
-    pipeline_parser = subparsers.add_parser('pipeline', help='Full pipeline')
-    pipeline_parser.add_argument('--task', '-t', required=True, nargs='+', choices=VALID_TASKS)
-    pipeline_parser.add_argument('--verify', action='store_true')
-    pipeline_parser.add_argument('--skip-verify', '-S', action='store_true', help='Skip verification in download stage')
-    
-    # List tasks
-    subparsers.add_parser('list-tasks', help='List available tasks')
-    
-    args = parser.parse_args()
-    
-    if args.command is None:
-        parser.print_help()
-        return
-    
-    commands = {
-        'init': init_project,
-        'download': run_download,
-        'download-single': run_download_single,
-        'verify': run_verify,
-        'pipeline': run_pipeline,
-        'list-tasks': list_tasks,
-    }
-    
-    commands[args.command](args)
+
+    # 완료 메시지
+    total = sum(results.values())
+    print(f"\n{'='*60}")
+    print(f"완료! 총 {total}개 저장")
+    for cat, cnt in results.items():
+        print(f"  - {CATEGORY_NAMES[cat]}: {cnt}개")
+    print(f"{'='*60}\n")
+
+    return 0
 
 
-if __name__ == '__main__':
-    main()
+if __name__ == "__main__":
+    sys.exit(main())

ytcollector/config.py

@@ -1,67 +1,85 @@
-"""
-SBS Dataset Collector - Configuration
-"""
-from pathlib import Path
-import platform
+# 설정 상수
 
-# Default paths (will use current working directory)
-def get_paths(base_dir: Path = None):
-    """Get all paths based on base directory"""
-    if base_dir is None:
-        base_dir = Path.cwd()
-    
-    return {
-        'base': base_dir,
-        'data': base_dir / "data",
-        # 'urls' removed - now inside video/{task}/youtube_url.txt
-        'videos': base_dir / "data" / "videos",  # 원본 전체 영상
-        'clips': base_dir / "video",             # 클리핑된 영상 (요구사항: video/task_이름)
-        'outputs': base_dir / "outputs",
-        'reports': base_dir / "outputs" / "reports",
-        'history': base_dir / "download_history.json",
-    }
+USER_AGENTS = [
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:121.0) Gecko/20100101 Firefox/121.0',
+    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.2 Safari/605.1.15',
+]
 
-# 사용자 지정 출력 경로 (네트워크 드라이브 등)
-# macOS에서는 "/Volumes/Data/..." 등으로 마운트된 경로를 사용해야 함
-NAS_PATH_WINDOWS = r"\\NAS_SERVER_IP\Data\Private Dataset\SBS_De-Identification_YouTube"
-NAS_PATH_MAC = "/Volumes/Data/Private Dataset/SBS_De-Identification_YouTube"
-
-if platform.system() == 'Windows':
-    CUSTOM_OUTPUT_DIR = NAS_PATH_WINDOWS
-elif platform.system() == 'Darwin': # macOS
-    CUSTOM_OUTPUT_DIR = NAS_PATH_MAC
-else:
-    CUSTOM_OUTPUT_DIR = None
-
-# Video settings
-CLIP_DURATION_BEFORE = 90  # 1분 30초 (초 단위)
-CLIP_DURATION_AFTER = 90   # 1분 30초 (초 단위)
-MAX_CLIP_DURATION = 180    # 최대 3분
-
-# Download settings
-VIDEO_FORMAT = "best[ext=mp4]/best"
-DOWNLOAD_RETRIES = 3
-MAX_VIDEOS_PER_TASK = 1000  # 태스크별 최대 영상 저장 수 (CLI -n 옵션으로 덮어쓰기 가능)
-
-# Fast Mode Settings (Parallel)
-MAX_WORKERS = 4              # 병렬 작업 프로세스 수
-REQUEST_DELAY_MIN = 1.0      # 최소 대기 시간 (초)
-REQUEST_DELAY_MAX = 3.0      # 최대 대기 시간 (초)
-PROXY_URL = None             # 프록시 (예: "http://user:pass@host:port")
+# 카테고리별 검색어 - SBS 콘텐츠 중심
+CATEGORY_QUERIES = {
+    'face': [
+        "SBS 인터뷰 클립",
+        "런닝맨 멤버 인터뷰",
+        "SBS 뉴스 인터뷰",
+        "미운우리새끼 인터뷰",
+        "SBS 스페셜 인물",
+        "집사부일체 인터뷰",
+        "그것이알고싶다 인터뷰",
+        "SBS 연예대상 소감",
+    ],
+    'license_plate': [
+        "중고차 매물 소개",
+        "자동차 세차 영상",
+        "신차 출고 브이로그",
+        "자동차 튜닝 작업",
+        "엔카 허위매물",
+        "주차장 만차",
+    ],
+    'tattoo': [
+        "타투 시술 영상",
+        "tattoo timelapse",
+        "타투이스트 작업",
+        "tattoo artist work",
+        "문신 시술",
+        "tattoo session",
+    ],
+    'text': [
+        "SBS 런닝맨 레전드",
+        "SBS 미운우리새끼 명장면",
+        "SBS 동상이몽 클립",
+        "SBS 집사부일체 모음",
+        "SBS 골목식당 레전드",
+        "SBS 맛남의광장 클립",
+        "SBS 불타는청춘 명장면",
+        "SBS 정글의법칙 레전드",
+        "SBS 예능",
+    ],
+}
 
-# YOLO-World settings
-YOLO_MODEL = "yolov8s-worldv2.pt"
-CONFIDENCE_THRESHOLD = 0.25
-FRAME_SAMPLE_RATE = 30  # 매 30프레임마다 샘플링 (약 1초)
+CATEGORY_NAMES = {
+    'face': '얼굴',
+    'license_plate': '번호판',
+    'tattoo': '타투',
+    'text': '텍스트'
+}
 
-# Task-specific class prompts
-TASK_CLASSES = {
-    "face": ["face"],
-    "license_plate": ["license plate"],
-    "tattoo": ["tattoo"],
-    "text": ["text"],
-    "knife": ["knife"],
-    "cigarette": ["cigarette"]
+# 카테고리별 제외 키워드 (제목에 포함 시 스킵)
+BLACKLIST_KEYWORDS = {
+    'tattoo': [
+        "두피 문신", "두피문신",
+        "눈썹 문신", "눈썹문신",
+        "입술 문신", "입술문신",
+        "틴트 입술",
+        "반영구", "SMP"
+    ],
+    'face': [],
+    'license_plate': [],
+    'text': []
 }
 
-VALID_TASKS = list(TASK_CLASSES.keys())
+# 번호판 정규식 패턴
+LICENSE_PLATE_PATTERNS = [
+    r'\d{2,3}[가-힣]\d{4}',
+    r'[가-힣]{2}\d{2}[가-힣]\d{4}',
+    r'[A-Z]{2,3}[-\s]?\d{2,4}[-\s]?[A-Z]{0,3}',
+    r'\d{2,4}[-\s]?[A-Z]{2,3}[-\s]?\d{2,4}',
+]
+
+# 스킵할 에러 메시지
+SKIP_ERRORS = [
+    "not available", "unavailable", "private", "removed",
+    "deleted", "copyright", "blocked", "age", "sign in",
+    "members-only", "premiere", "live event"
+]

ytcollector/dataset_builder.py

@@ -0,0 +1,71 @@
+import os
+import subprocess
+from yt_dlp import YoutubeDL
+from .utils import clip_video, get_url_list, get_video_duration, timestamp_to_seconds
+
+def download_videos(url_list, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+    for idx, item in enumerate(url_list, 1):
+        url = item['url']
+        task = item['task']
+        index_str = f"{idx:03d}"
+        
+        existing_files = [f for f in os.listdir(output_dir) if f.startswith(f"{index_str}_")]
+        if existing_files:
+            print(f"[{index_str}] Skip: {existing_files[0]}")
+            continue
+
+        print(f"[{index_str}] Downloading: {url} ({task})")
+        try:
+            ydl_opts = {
+                'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
+                'outtmpl': os.path.join(output_dir, f"{index_str}_{task}_%(title)s.%(ext)s"),
+                'quiet': True,
+                'no_warnings': True,
+            }
+            with YoutubeDL(ydl_opts) as ydl:
+                ydl.download([url])
+        except Exception as e:
+            print(f"[{index_str}] Failed: {e}")
+
+def build_dataset(url_file, output_root="."):
+    video_dir = os.path.abspath(os.path.join(output_root, "video"))
+    clip_dir = os.path.abspath(os.path.join(output_root, "video_clips"))
+    
+    urls = get_url_list(url_file)
+    if not urls:
+        print(f"Error: No valid data in {url_file}")
+        return
+
+    print(f"--- Step 1: Downloading {len(urls)} videos ---")
+    download_videos(urls, video_dir)
+    
+    print(f"\n--- Step 2: Clipping videos ---")
+    os.makedirs(clip_dir, exist_ok=True)
+    for idx, item in enumerate(urls, 1):
+        index_str = f"{idx:03d}"
+        files = [f for f in os.listdir(video_dir) if f.startswith(f"{index_str}_")]
+        if not files: continue
+            
+        input_file = os.path.join(video_dir, files[0])
+        output_file = os.path.join(clip_dir, files[0])
+        
+        if os.path.exists(output_file): continue
+            
+        print(f"[{index_str}] Clipping: {files[0]}")
+        center_sec = timestamp_to_seconds(item['timestamp'])
+        clip_video(input_file, output_file, center_sec)
+
+    print(f"\nDone! Clips saved in: {clip_dir}")
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description='Build SBS Dataset from YouTube URL list')
+    parser.add_argument('file', help='Path to youtube_url.txt')
+    parser.add_argument('-o', '--output', default='.', help='Output root directory (default: .)')
+    args = parser.parse_args()
+    
+    build_dataset(args.file, args.output)
+
+if __name__ == "__main__":
+    main()