MogensR commited on
Commit
968e618
·
1 Parent(s): dc14288

Update processing/video/video_processor.py

Browse files
Files changed (1) hide show
  1. processing/video/video_processor.py +105 -63
processing/video/video_processor.py CHANGED
@@ -2,105 +2,144 @@
2
  """
3
  Compatibility shim: CoreVideoProcessor
4
 
5
- Bridges the legacy import `from processing.video.video_processor import CoreVideoProcessor`
6
- to the modern pipeline functions living in `utils.cv_processing` and models in `core.models`.
 
 
 
 
 
 
7
  """
8
 
9
  from __future__ import annotations
10
 
11
  from dataclasses import dataclass
12
  from typing import Optional, Dict, Any, Tuple, Callable
 
 
13
 
14
  import cv2
15
  import numpy as np
16
- import time
17
- import threading
18
 
19
- from utils.logger import get_logger
20
- from core.models import ModelManager
 
 
 
 
 
 
21
 
22
- # these funcs are the ones you showed (in utils/cv_processing.py)
23
  from utils.cv_processing import (
24
  segment_person_hq,
25
  refine_mask_hq,
26
  replace_background_hq,
27
  create_professional_background,
28
  validate_video_file,
 
29
  )
30
 
 
31
  @dataclass
32
  class ProcessorConfig:
33
  background_preset: str = "minimalist" # key in PROFESSIONAL_BACKGROUNDS
34
  write_fps: Optional[float] = None # None -> keep source fps
35
 
 
36
  class CoreVideoProcessor:
37
  """
38
  Minimal, safe implementation used by core/app.py.
39
- It relies on ModelManager (SAM2 + MatAnyone) and your cv_processing helpers.
40
- Now supports live progress + cancel/stop.
 
 
 
 
41
  """
42
 
43
- def __init__(self, config: Optional[ProcessorConfig] = None, models: Optional[ModelManager] = None):
44
- self.log = get_logger(f"{__name__}.CoreVideoProcessor")
45
  self.config = config or ProcessorConfig()
46
- self.models = models or ModelManager()
47
- try:
48
- self.models.load_all()
49
- except Exception as e:
50
- self.log.warning(f"Model load issue (will use fallbacks if needed): {e}")
51
 
52
- # --- single-frame API (useful for images or per-frame video loop) ---
53
  def process_frame(self, frame: np.ndarray, background: np.ndarray) -> Dict[str, Any]:
54
- """Return dict with composited frame + mask; always succeeds with fallbacks."""
55
  predictor = None
56
  try:
57
- sam2_model = self.models.get_sam2()
58
- if sam2_model is not None:
59
- if hasattr(sam2_model, 'predictor'):
60
- predictor = sam2_model.predictor
61
- elif hasattr(sam2_model, 'set_image'):
62
- predictor = sam2_model
63
- elif isinstance(sam2_model, dict) and 'model' in sam2_model:
64
- self.log.warning("SAM2 loaded as dict format, not directly usable")
65
- predictor = None
66
- if predictor is None:
67
- self.log.debug("SAM2 predictor not available, will use fallback")
68
  except Exception as e:
69
  self.log.warning(f"SAM2 predictor unavailable: {e}")
70
 
71
- # 1) segment
72
  mask = segment_person_hq(frame, predictor, fallback_enabled=True)
73
 
74
- # 2) refine
75
  matanyone = None
76
  try:
77
- matanyone_model = self.models.get_matanyone()
78
- if matanyone_model is not None:
79
- matanyone = matanyone_model
80
  except Exception as e:
81
- self.log.warning(f"MatAnyone unavailable: {e}")
82
 
83
  mask_refined = refine_mask_hq(frame, mask, matanyone, fallback_enabled=True)
84
 
85
- # 3) composite
86
  out = replace_background_hq(frame, mask_refined, background, fallback_enabled=True)
87
 
88
  return {"frame": out, "mask": mask_refined}
89
 
90
- # --- simple video API (covers typical usage in older core/app.py code) ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91
  def process_video(
92
  self,
93
  input_path: str,
94
  output_path: str,
95
  bg_config: Optional[Dict[str, Any]] = None,
96
- progress_callback: Optional[Callable[[int, int, float], None]] = None, # <-- ADDED
97
- stop_event: Optional[threading.Event] = None # <-- ADDED
98
  ) -> Dict[str, Any]:
99
  """
100
- Process a full video with live progress and optional stop.
101
- progress_callback: function(current_frame, total_frames, fps)
102
- stop_event: threading.Event() - if set(), abort processing.
103
- Returns: dict with stats.
104
  """
105
  ok, msg = validate_video_file(input_path)
106
  if not ok:
@@ -115,16 +154,16 @@ def process_video(
115
  height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
116
  fps = cap.get(cv2.CAP_PROP_FPS)
117
  total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
118
- fps_out = self.config.write_fps or (fps if fps and fps > 0 else 25.0)
119
 
 
120
  fourcc = cv2.VideoWriter_fourcc(*"mp4v")
121
- writer = cv2.VideoWriter(output_path, fourcc, fps_out, (width, height))
 
 
 
122
 
123
- # Build background (once)
124
- from utils.cv_processing import PROFESSIONAL_BACKGROUNDS
125
- preset = self.config.background_preset
126
- cfg = bg_config or PROFESSIONAL_BACKGROUNDS.get(preset, PROFESSIONAL_BACKGROUNDS["minimalist"])
127
- background = create_professional_background(cfg, width, height)
128
 
129
  frame_count = 0
130
  start_time = time.time()
@@ -134,24 +173,25 @@ def process_video(
134
  if not ret:
135
  break
136
 
137
- # --- CANCEL SUPPORT ---
138
  if stop_event is not None and stop_event.is_set():
139
- self.log.info("Processing stopped by user request") # <-- CHANGED
140
  break
141
 
142
- res = self.process_frame(frame, background)
143
- writer.write(res["frame"])
 
144
  frame_count += 1
145
 
146
- # --- LIVE PROGRESS ---
147
  if progress_callback:
148
  elapsed = time.time() - start_time
149
- fps_live = frame_count / elapsed if elapsed > 0 else 0
150
- progress_callback(
151
- frame_count,
152
- total_frames,
153
- fps_live
154
- )
155
  finally:
156
  cap.release()
157
  writer.release()
@@ -161,8 +201,10 @@ def process_video(
161
  "frames": frame_count,
162
  "width": width,
163
  "height": height,
164
- "fps_out": fps_out
 
165
  }
166
 
167
- # Backward-compat export name
 
168
  VideoProcessor = CoreVideoProcessor
 
2
  """
3
  Compatibility shim: CoreVideoProcessor
4
 
5
+ Bridges the legacy import
6
+ from processing.video.video_processor import CoreVideoProcessor
7
+ to the modern pipeline functions in utils.cv_processing, using whatever
8
+ models provider is passed in (e.g., models.loaders.ModelLoader).
9
+
10
+ Requirements for the models provider:
11
+ - get_sam2() -> predictor or None
12
+ - get_matanyone() -> processor or None
13
  """
14
 
15
  from __future__ import annotations
16
 
17
  from dataclasses import dataclass
18
  from typing import Optional, Dict, Any, Tuple, Callable
19
+ import time
20
+ import threading
21
 
22
  import cv2
23
  import numpy as np
 
 
24
 
25
+ # Try project logger; fall back to std logging
26
+ try:
27
+ from utils.logger import get_logger
28
+ _log = get_logger(__name__)
29
+ except Exception:
30
+ import logging
31
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
32
+ _log = logging.getLogger(__name__)
33
 
34
+ # CV pipeline helpers
35
  from utils.cv_processing import (
36
  segment_person_hq,
37
  refine_mask_hq,
38
  replace_background_hq,
39
  create_professional_background,
40
  validate_video_file,
41
+ PROFESSIONAL_BACKGROUNDS,
42
  )
43
 
44
+
45
  @dataclass
46
  class ProcessorConfig:
47
  background_preset: str = "minimalist" # key in PROFESSIONAL_BACKGROUNDS
48
  write_fps: Optional[float] = None # None -> keep source fps
49
 
50
+
51
  class CoreVideoProcessor:
52
  """
53
  Minimal, safe implementation used by core/app.py.
54
+ It relies on a models provider (e.g., ModelLoader) that implements:
55
+ - get_sam2()
56
+ - get_matanyone()
57
+ and uses utils.cv_processing for the pipeline.
58
+
59
+ Supports progress callback and cancellation via stop_event.
60
  """
61
 
62
+ def __init__(self, config: Optional[ProcessorConfig] = None, models: Optional[Any] = None):
63
+ self.log = _log
64
  self.config = config or ProcessorConfig()
65
+ self.models = models # do NOT load here; core/app handles loading
66
+ if self.models is None:
67
+ self.log.warning("CoreVideoProcessor initialized without a models provider; will use fallbacks.")
 
 
68
 
69
+ # ---------- Single frame ----------
70
  def process_frame(self, frame: np.ndarray, background: np.ndarray) -> Dict[str, Any]:
71
+ """Return dict with composited frame + mask; always attempts fallbacks."""
72
  predictor = None
73
  try:
74
+ if self.models and hasattr(self.models, "get_sam2"):
75
+ predictor = self.models.get_sam2()
76
+ # Some wrappers expose predictor directly, others are already usable
77
+ # segment_person_hq checks for set_image/predict itself.
 
 
 
 
 
 
 
78
  except Exception as e:
79
  self.log.warning(f"SAM2 predictor unavailable: {e}")
80
 
81
+ # 1) segmentation (with fallbacks inside)
82
  mask = segment_person_hq(frame, predictor, fallback_enabled=True)
83
 
84
+ # 2) refinement (MatAnyOne if available, else robust OpenCV path)
85
  matanyone = None
86
  try:
87
+ if self.models and hasattr(self.models, "get_matanyone"):
88
+ matanyone = self.models.get_matanyone()
 
89
  except Exception as e:
90
+ self.log.warning(f"MatAnyOne unavailable: {e}")
91
 
92
  mask_refined = refine_mask_hq(frame, mask, matanyone, fallback_enabled=True)
93
 
94
+ # 3) compositing
95
  out = replace_background_hq(frame, mask_refined, background, fallback_enabled=True)
96
 
97
  return {"frame": out, "mask": mask_refined}
98
 
99
+ # ---------- Build background once per video ----------
100
+ def _prepare_background_from_config(
101
+ self,
102
+ bg_config: Optional[Dict[str, Any]],
103
+ width: int,
104
+ height: int
105
+ ) -> np.ndarray:
106
+ """
107
+ Accepts either:
108
+ - {"custom_path": "/path/to/image.png"} → load image
109
+ - {"background_choice": "minimalist"} → preset
110
+ - None → use self.config.background_preset
111
+ """
112
+ # 1) custom image?
113
+ if bg_config and bg_config.get("custom_path"):
114
+ path = bg_config["custom_path"]
115
+ img = cv2.imread(path, cv2.IMREAD_COLOR)
116
+ if img is None:
117
+ self.log.warning(f"Custom background at '{path}' could not be read. Falling back to preset.")
118
+ else:
119
+ return cv2.resize(img, (width, height), interpolation=cv2.INTER_LANCZOS4)
120
+
121
+ # 2) preset (explicit choice or default)
122
+ choice = None
123
+ if bg_config and "background_choice" in bg_config:
124
+ choice = bg_config["background_choice"]
125
+ if not choice:
126
+ choice = self.config.background_preset
127
+
128
+ cfg = PROFESSIONAL_BACKGROUNDS.get(choice, PROFESSIONAL_BACKGROUNDS["minimalist"])
129
+ return create_professional_background(cfg, width, height)
130
+
131
+ # ---------- Full video ----------
132
  def process_video(
133
  self,
134
  input_path: str,
135
  output_path: str,
136
  bg_config: Optional[Dict[str, Any]] = None,
137
+ progress_callback: Optional[Callable[[int, int, float], None]] = None,
138
+ stop_event: Optional[threading.Event] = None
139
  ) -> Dict[str, Any]:
140
  """
141
+ Process a full video with live progress and optional cancel.
142
+ progress_callback(current_frame, total_frames, fps_live)
 
 
143
  """
144
  ok, msg = validate_video_file(input_path)
145
  if not ok:
 
154
  height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
155
  fps = cap.get(cv2.CAP_PROP_FPS)
156
  total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
157
 
158
+ fps_out = self.config.write_fps or (fps if fps and fps > 0 else 25.0)
159
  fourcc = cv2.VideoWriter_fourcc(*"mp4v")
160
+ writer = cv2.VideoWriter(output_path, fourcc, float(fps_out), (width, height))
161
+ if not writer.isOpened():
162
+ cap.release()
163
+ raise RuntimeError(f"Could not open writer for: {output_path}")
164
 
165
+ # Build background once
166
+ background = self._prepare_background_from_config(bg_config, width, height)
 
 
 
167
 
168
  frame_count = 0
169
  start_time = time.time()
 
173
  if not ret:
174
  break
175
 
176
+ # Cancel support
177
  if stop_event is not None and stop_event.is_set():
178
+ self.log.info("Processing stopped by user request.")
179
  break
180
 
181
+ # Process single frame
182
+ result = self.process_frame(frame, background)
183
+ writer.write(result["frame"])
184
  frame_count += 1
185
 
186
+ # Progress callback
187
  if progress_callback:
188
  elapsed = time.time() - start_time
189
+ fps_live = frame_count / elapsed if elapsed > 0 else 0.0
190
+ try:
191
+ progress_callback(frame_count, total_frames, fps_live)
192
+ except Exception:
193
+ # Don’t break processing due to a UI callback error
194
+ pass
195
  finally:
196
  cap.release()
197
  writer.release()
 
201
  "frames": frame_count,
202
  "width": width,
203
  "height": height,
204
+ "fps_out": float(fps_out),
205
+ "output_path": output_path,
206
  }
207
 
208
+
209
+ # Backward-compat alias used elsewhere
210
  VideoProcessor = CoreVideoProcessor