File size: 3,738 Bytes
ad2cb5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import json
import subprocess
from pathlib import Path
from config import make_path
from files.utils.logging import get_logger


class FrameExtractor:
    def __init__(self, video_path: str, min_scene_len: float = 0.2):
        self.min_scene_len = min_scene_len
        self.video_path = Path(video_path)
        self.scene_json_path = self.frame_json = make_path('processed/scene-detection', video_path, 'scene', 'json')
        self.output_dir = make_path('interim/frames', video_path, '', '')
        self.output_dir.mkdir(parents=True, exist_ok=True)

        log_file = f'{self.video_path.stem}_log.txt'
        self.logger = get_logger('frame_extract', log_file)

    def _ffmpeg_extract(self, timestamp: float, out_path: Path):
        cmd = [
            'ffmpeg',
            '-loglevel', 'error',
            '-y',
            '-ss', f'{timestamp:.3f}',
            '-t', '1',
            '-i', str(self.video_path),
            '-frames:v', '1',
            '-q:v', '2',
            '-pix_fmt', 'yuvj420p',
            str(out_path)
        ]
        result = subprocess.run(cmd, capture_output=True)
        if result.returncode != 0:
            self.logger.error('ffmpeg failed: %s', result.stderr.decode('utf-8', 'ignore').strip())

    def _get_brightness(self, timestamp: float) -> float:
        cmd = [
            'ffprobe',
            '-v', 'error',
            '-read_intervals', f'%{timestamp}+1',
            '-select_streams', 'v:0',
            '-show_frames',
            '-show_entries', 'frame_tags=lavfi.signalstats.YAVG',
            '-of', 'default=noprint_wrappers=1:nokey=1',
            str(self.video_path)
        ]
        result = subprocess.run(cmd, capture_output=True, text=True)
        try:
            yavg_values = [float(line.strip()) for line in result.stdout.strip().split('\n') if line.strip()]
            if yavg_values:
                return yavg_values[0]
        except Exception:
            pass
        self.logger.warning('Could not get brightness at %.2fs', timestamp)
        return -1.0

    def extract(self) -> list[dict]:
        with open(self.scene_json_path, encoding='utf-8') as f:
            scenes = json.load(f).get('scenes', [])
        if not scenes:
            self.logger.warning('No scenes found in %s', self.scene_json_path)
            return []

        delta = 0.5
        results = []

        for i, sc in enumerate(scenes):
            start = float(sc['start_time'])
            end = float(sc['end_time'])
            dur = end - start
            if dur < self.min_scene_len:
                self.logger.warning('Scene %s too short (%.2fs), skipping', i, dur)
                continue

            mid = (start + end) / 2

            frame_path = self.output_dir / f'{self.video_path.stem}_scene_{i:02}.jpg'
            prev_path = self.output_dir / f'{self.video_path.stem}_scene_{i:02}_prev.jpg'
            next_path = self.output_dir / f'{self.video_path.stem}_scene_{i:02}_next.jpg'

            self._ffmpeg_extract(mid, frame_path)
            self._ffmpeg_extract(mid - delta, prev_path)
            self._ffmpeg_extract(mid + delta, next_path)

            brightness = self._get_brightness(mid)

            self.logger.info('[Scene %s] %.2fs → %s | Brightness: %.2f', i, mid, frame_path.name, brightness)

            results.append({
                'scene_index': i,
                'timestamp': mid,
                'frame_path': str(frame_path),
                'prev_frame_path': str(prev_path),
                'next_frame_path': str(next_path),
                'brightness': brightness
            })

        self.logger.info('%s frames (with context) extracted to %s', len(results), self.output_dir)
        return results