jbilcke-hf HF staff commited on
Commit
a529bb7
·
1 Parent(s): 1b19314

add automatic FPS tracking

Browse files
Files changed (2) hide show
  1. finetrainers_utils.py +9 -1
  2. utils.py +42 -0
finetrainers_utils.py CHANGED
@@ -4,7 +4,7 @@ import logging
4
  import shutil
5
  from typing import Any, Optional, Dict, List, Union, Tuple
6
  from config import STORAGE_PATH, TRAINING_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, HF_API_TOKEN, MODEL_TYPES
7
- from utils import extract_scene_info, make_archive, is_image_file, is_video_file
8
 
9
  logger = logging.getLogger(__name__)
10
 
@@ -110,6 +110,14 @@ def copy_files_to_training_dir(prompt_prefix: str) -> int:
110
  if parent_caption and not caption.endswith(parent_caption):
111
  caption = f"{caption}\n{parent_caption}"
112
 
 
 
 
 
 
 
 
 
113
  if prompt_prefix and not caption.startswith(prompt_prefix):
114
  caption = f"{prompt_prefix}{caption}"
115
 
 
4
  import shutil
5
  from typing import Any, Optional, Dict, List, Union, Tuple
6
  from config import STORAGE_PATH, TRAINING_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, HF_API_TOKEN, MODEL_TYPES
7
+ from utils import get_video_fps, extract_scene_info, make_archive, is_image_file, is_video_file
8
 
9
  logger = logging.getLogger(__name__)
10
 
 
110
  if parent_caption and not caption.endswith(parent_caption):
111
  caption = f"{caption}\n{parent_caption}"
112
 
113
+ # Add FPS information for videos
114
+ if is_video_file(file_path) and caption:
115
+ # Only add FPS if not already present
116
+ if not any(f"FPS, " in line for line in caption.split('\n')):
117
+ fps_info = get_video_fps(file_path)
118
+ if fps_info:
119
+ caption = f"{fps_info}{caption}"
120
+
121
  if prompt_prefix and not caption.startswith(prompt_prefix):
122
  caption = f"{prompt_prefix}{caption}"
123
 
utils.py CHANGED
@@ -1,11 +1,15 @@
1
  import os
2
  import shutil
 
3
  from huggingface_hub import HfApi, create_repo
4
  from pathlib import Path
5
  import json
6
  import re
 
7
  from typing import Any, Optional, Dict, List, Union, Tuple
8
 
 
 
9
  def make_archive(source: str | Path, destination: str | Path):
10
  source = str(source)
11
  destination = str(destination)
@@ -18,6 +22,44 @@ def make_archive(source: str | Path, destination: str | Path):
18
  shutil.make_archive(name, format, archive_from, archive_to)
19
  shutil.move('%s.%s'%(name,format), destination)
20
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21
  def extract_scene_info(filename: str) -> Tuple[str, Optional[int]]:
22
  """Extract base name and scene number from filename
23
 
 
1
  import os
2
  import shutil
3
+ import subprocess
4
  from huggingface_hub import HfApi, create_repo
5
  from pathlib import Path
6
  import json
7
  import re
8
+ import logging
9
  from typing import Any, Optional, Dict, List, Union, Tuple
10
 
11
+ logger = logging.getLogger(__name__)
12
+
13
  def make_archive(source: str | Path, destination: str | Path):
14
  source = str(source)
15
  destination = str(destination)
 
22
  shutil.make_archive(name, format, archive_from, archive_to)
23
  shutil.move('%s.%s'%(name,format), destination)
24
 
25
+ def get_video_fps(video_path: Path) -> Optional[str]:
26
+ """Get FPS information from video file using ffprobe
27
+
28
+ Args:
29
+ video_path: Path to video file
30
+
31
+ Returns:
32
+ FPS string (e.g. "24 FPS, ") or None if unable to determine
33
+ """
34
+ try:
35
+ cmd = [
36
+ 'ffprobe',
37
+ '-v', 'error',
38
+ '-select_streams', 'v:0',
39
+ '-show_entries', 'stream=avg_frame_rate',
40
+ '-of', 'default=noprint_wrappers=1:nokey=1',
41
+ str(video_path)
42
+ ]
43
+
44
+ result = subprocess.run(cmd, capture_output=True, text=True)
45
+ if result.returncode != 0:
46
+ logger.warning(f"Error getting FPS for {video_path}: {result.stderr}")
47
+ return None
48
+
49
+ fps = result.stdout.strip()
50
+ if '/' in fps:
51
+ # Convert fraction to decimal
52
+ num, den = map(int, fps.split('/'))
53
+ if den == 0:
54
+ return None
55
+ fps = str(round(num / den))
56
+
57
+ return f"{fps} FPS, "
58
+
59
+ except Exception as e:
60
+ logger.warning(f"Failed to get FPS for {video_path}: {e}")
61
+ return None
62
+
63
  def extract_scene_info(filename: str) -> Tuple[str, Optional[int]]:
64
  """Extract base name and scene number from filename
65