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

ytcollector/__init__.py

@@ -1,14 +1,39 @@
-"""
-SBS Dataset Collector - YouTube 영상 수집 및 YOLO-World 검증 파이프라인
-"""
-from pathlib import Path
+"""YouTube 콘텐츠 수집기 라이브러리
 
-__version__ = "1.0.8"
-__author__ = "SBS Dataset Team"
+외부에서 라이브러리로 사용하거나 CLI로 실행할 수 있습니다.
+
+라이브러리 사용 예시:
+    from ytcollector import YouTubeDownloader, run
+    
+    # 방법 1: YouTubeDownloader 직접 사용
+    downloader = YouTubeDownloader(output_path="./videos")
+    count = downloader.collect("face", max_videos=5)
+    
+    # 방법 2: run() 함수 사용 (간단한 방법)
+    results = run(categories=["face", "text"], count=3)
+    
+CLI 사용 예시:
+    ytcollector -c face -n 5
+    ytc -c face text --fast
+"""
 
-# Package root directory
-PACKAGE_DIR = Path(__file__).parent
+from .config import CATEGORY_NAMES, CATEGORY_QUERIES, USER_AGENTS, LICENSE_PLATE_PATTERNS
+from .analyzer import VideoAnalyzer, check_dependencies
+from .downloader import YouTubeDownloader
+from .cli import run, main as cli_main
 
-# Default data directories (can be overridden)
-DEFAULT_DATA_DIR = Path.cwd() / "data"
-DEFAULT_OUTPUT_DIR = Path.cwd() / "outputs"
+__version__ = "1.0.0"
+__all__ = [
+    # 주요 클래스
+    "VideoAnalyzer",
+    "YouTubeDownloader",
+    # 설정
+    "CATEGORY_NAMES",
+    "CATEGORY_QUERIES",
+    "USER_AGENTS",
+    "LICENSE_PLATE_PATTERNS",
+    # 유틸리티
+    "check_dependencies",
+    "run",
+    "cli_main",
+]

ytcollector/analyzer.py

@@ -0,0 +1,239 @@
+import re
+import threading
+from .config import LICENSE_PLATE_PATTERNS
+
+# 선택적 import
+try:
+    import cv2
+    CV2_AVAILABLE = True
+except ImportError:
+    CV2_AVAILABLE = False
+
+try:
+    import easyocr
+    EASYOCR_AVAILABLE = True
+except ImportError:
+    EASYOCR_AVAILABLE = False
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+
+
+class VideoAnalyzer:
+    """영상 분석 클래스 - 얼굴, 텍스트, 번호판, 타투 감지"""
+
+    _ocr_lock = threading.Lock()
+
+    def __init__(self):
+        self.ocr_reader = None
+        self.face_cascade = None
+
+        if CV2_AVAILABLE:
+            cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
+            self.face_cascade = cv2.CascadeClassifier(cascade_path)
+
+    def _init_ocr(self):
+        """OCR 리더 초기화 (필요할 때만, 스레드 안전)"""
+        if EASYOCR_AVAILABLE and self.ocr_reader is None:
+            with self._ocr_lock:
+                if self.ocr_reader is None:
+                    print("  OCR 엔진 초기화 중...")
+                    self.ocr_reader = easyocr.Reader(['ko', 'en'], gpu=False, verbose=False)
+
+    def extract_frames(self, video_path, num_frames=10):
+        """영상에서 균등 간격으로 프레임 추출"""
+        if not CV2_AVAILABLE:
+            return []
+
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            return []
+
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        if total_frames <= 0:
+            cap.release()
+            return []
+
+        frame_indices = [int(i * total_frames / (num_frames + 1)) for i in range(1, num_frames + 1)]
+        frames = []
+
+        for idx in frame_indices:
+            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
+            ret, frame = cap.read()
+            if ret:
+                frames.append(frame)
+
+        cap.release()
+        return frames
+
+    def detect_faces(self, frame):
+        """Haar Cascade로 얼굴 감지"""
+        if not CV2_AVAILABLE or self.face_cascade is None:
+            return []
+
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        return self.face_cascade.detectMultiScale(
+            gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
+        )
+
+    def detect_text(self, frame):
+        """EasyOCR로 텍스트 감지 (스레드 안전)"""
+        if not EASYOCR_AVAILABLE:
+            return []
+
+        self._init_ocr()
+        try:
+            h, w = frame.shape[:2]
+            if w > 640:
+                scale = 640 / w
+                frame = cv2.resize(frame, (640, int(h * scale)))
+
+            with self._ocr_lock:
+                results = self.ocr_reader.readtext(frame)
+            return [r[1] for r in results if r[2] > 0.3]
+        except:
+            return []
+
+    def detect_license_plate(self, texts):
+        """텍스트에서 번호판 패턴 감지"""
+        for text in texts:
+            text_clean = text.replace(' ', '').upper()
+            for pattern in LICENSE_PLATE_PATTERNS:
+                if re.search(pattern, text_clean):
+                    return True
+        return False
+
+    def detect_tattoo(self, frame):
+        """피부 영역에서 타투(어두운 잉크 패턴) 감지"""
+        if not CV2_AVAILABLE or not NUMPY_AVAILABLE:
+            return False
+
+        try:
+            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+
+            # 피부색 범위
+            lower_skin = np.array([0, 30, 80], dtype=np.uint8)
+            upper_skin = np.array([17, 170, 255], dtype=np.uint8)
+            skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
+
+            # 노이즈 제거
+            kernel = np.ones((5, 5), np.uint8)
+            skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
+            skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
+
+            skin_pixels = cv2.countNonZero(skin_mask)
+            total_pixels = frame.shape[0] * frame.shape[1]
+
+            # 피부 영역 최소 10% 필요
+            if skin_pixels < total_pixels * 0.10:
+                return False
+
+            # 피부 영역 내 어두운 픽셀(타투) 감지
+            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+            skin_gray = cv2.bitwise_and(gray, gray, mask=skin_mask)
+            dark_mask = cv2.inRange(skin_gray, 1, 80)
+
+            dark_pixels = cv2.countNonZero(dark_mask)
+            dark_ratio = dark_pixels / max(skin_pixels, 1)
+
+            # 어두운 영역이 3~35%일 때 타투로 판정
+            if 0.03 < dark_ratio < 0.35:
+                contours, _ = cv2.findContours(dark_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                significant = [c for c in contours if cv2.contourArea(c) > 100]
+                return len(significant) >= 1
+
+            return False
+        except:
+            return False
+
+    def analyze(self, video_path):
+        """영상 전체 분석"""
+        results = {
+            'face': False,
+            'text': False,
+            'license_plate': False,
+            'tattoo': False,
+            'face_count': 0,
+            'detected_texts': [],
+            'first_detection_sec': None,
+            'first_detection_ts': None
+        }
+
+        if not CV2_AVAILABLE:
+            print("  ⚠ OpenCV 미설치")
+            return results
+
+        # 영상 정보 가져오기
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            return results
+        
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        cap.release()
+
+        num_analysis_frames = 10
+        frame_indices = [int(i * total_frames / (num_analysis_frames + 1)) for i in range(1, num_analysis_frames + 1)]
+        
+        all_texts = []
+        total_faces = 0
+
+        cap = cv2.VideoCapture(video_path)
+        for idx in frame_indices:
+            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
+            ret, frame = cap.read()
+            if not ret: continue
+
+            # 현재 프레임의 시간(초)
+            current_sec = idx / fps
+            detected_now = False
+
+            # 얼굴
+            faces = self.detect_faces(frame)
+            if len(faces) > 0:
+                results['face'] = True
+                total_faces += len(faces)
+                detected_now = True
+
+            # 텍스트
+            texts = self.detect_text(frame)
+            if texts:
+                results['text'] = True
+                all_texts.extend(texts)
+                detected_now = True
+
+            # 타투
+            if self.detect_tattoo(frame):
+                results['tattoo'] = True
+                detected_now = True
+
+            # 첫 감지 시점 기록
+            if detected_now and results['first_detection_sec'] is None:
+                results['first_detection_sec'] = current_sec
+                m, s = int(current_sec // 60), int(current_sec % 60)
+                results['first_detection_ts'] = f"{m:02d}:{s:02d}"
+
+        cap.release()
+
+        # 번호판 (텍스트에서)
+        if all_texts:
+            results['license_plate'] = self.detect_license_plate(all_texts)
+            results['detected_texts'] = list(set(all_texts))[:10]
+
+        results['face_count'] = total_faces
+        return results
+
+
+def check_dependencies():
+    """의존성 체크"""
+    missing = []
+    if not CV2_AVAILABLE:
+        missing.append("opencv-python")
+    if not EASYOCR_AVAILABLE:
+        missing.append("easyocr")
+    if not NUMPY_AVAILABLE:
+        missing.append("numpy")
+    return missing