MogensR commited on
Commit
0f29d58
·
1 Parent(s): 412ec39

Create utils__init__fixed.py

Browse files
Files changed (1) hide show
  1. utils__init__fixed.py +261 -0
utils__init__fixed.py ADDED
@@ -0,0 +1,261 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Complete utils/__init__.py with all required functions
3
+ Provides direct implementations to avoid import recursion
4
+ """
5
+
6
+ import cv2
7
+ import numpy as np
8
+ from PIL import Image
9
+ import torch
10
+ import logging
11
+ from typing import Optional, Tuple, Dict, Any, List
12
+ import tempfile
13
+ import os
14
+ from app.video_enhancer.matanyone_processor import MatAnyoneProcessor
15
+
16
+ logger = logging.getLogger(__name__)
17
+
18
+ # Cached MatAnyone processor (initialized on first use)
19
+ _MATANYONE_PROCESSOR: Optional[MatAnyoneProcessor] = None
20
+
21
+ # Professional backgrounds configuration
22
+ PROFESSIONAL_BACKGROUNDS = {
23
+ "office": {"color": (240, 248, 255), "gradient": True},
24
+ "studio": {"color": (32, 32, 32), "gradient": False},
25
+ "nature": {"color": (34, 139, 34), "gradient": True},
26
+ "abstract": {"color": (75, 0, 130), "gradient": True},
27
+ "white": {"color": (255, 255, 255), "gradient": False},
28
+ "black": {"color": (0, 0, 0), "gradient": False}
29
+ }
30
+
31
+ def validate_video_file(video_path: str) -> bool:
32
+ """Validate if video file is readable"""
33
+ try:
34
+ if not os.path.exists(video_path):
35
+ return False
36
+
37
+ cap = cv2.VideoCapture(video_path)
38
+ if not cap.isOpened():
39
+ return False
40
+
41
+ ret, frame = cap.read()
42
+ cap.release()
43
+ return ret and frame is not None
44
+
45
+ except Exception as e:
46
+ logger.error(f"Video validation failed: {e}")
47
+ return False
48
+
49
+ def segment_person_hq(frame: np.ndarray, use_sam2: bool = True) -> Optional[np.ndarray]:
50
+ """High-quality person segmentation using SAM2 or fallback methods"""
51
+ try:
52
+ if use_sam2:
53
+ # Try SAM2 segmentation
54
+ try:
55
+ from sam2.sam2_image_predictor import SAM2ImagePredictor
56
+ from sam2.build_sam import build_sam2
57
+ from huggingface_hub import hf_hub_download
58
+
59
+ # Load SAM2 model
60
+ sam_checkpoint = hf_hub_download("facebook/sam2-hiera-base-plus", "sam2_hiera_b+.pt")
61
+ sam_model = build_sam2(model_name='sam2_hiera_base_plus_t', ckpt_path=sam_checkpoint)
62
+ predictor = SAM2ImagePredictor(sam_model)
63
+
64
+ # Set image and predict
65
+ predictor.set_image(frame)
66
+
67
+ # Use center point as prompt (assuming person is in center)
68
+ h, w = frame.shape[:2]
69
+ center_point = np.array([[w//2, h//2]])
70
+ center_label = np.array([1])
71
+
72
+ masks, scores, _ = predictor.predict(
73
+ point_coords=center_point,
74
+ point_labels=center_label,
75
+ multimask_output=False
76
+ )
77
+
78
+ return masks[0] if len(masks) > 0 else None
79
+
80
+ except Exception as e:
81
+ logger.warning(f"SAM2 segmentation failed: {e}, falling back to simple method")
82
+
83
+ # Fallback: Simple person detection using background subtraction
84
+ return _simple_person_segmentation(frame)
85
+
86
+ except Exception as e:
87
+ logger.error(f"Person segmentation failed: {e}")
88
+ return None
89
+
90
+ def _simple_person_segmentation(frame: np.ndarray) -> np.ndarray:
91
+ """Simple person segmentation using color-based methods"""
92
+ # Convert to HSV for better color detection
93
+ hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
94
+
95
+ # Create mask for common background colors (green screen, white, etc.)
96
+ # Green screen detection
97
+ lower_green = np.array([40, 40, 40])
98
+ upper_green = np.array([80, 255, 255])
99
+ green_mask = cv2.inRange(hsv, lower_green, upper_green)
100
+
101
+ # White background detection
102
+ lower_white = np.array([0, 0, 200])
103
+ upper_white = np.array([180, 30, 255])
104
+ white_mask = cv2.inRange(hsv, lower_white, upper_white)
105
+
106
+ # Combine masks
107
+ bg_mask = cv2.bitwise_or(green_mask, white_mask)
108
+
109
+ # Invert to get person mask
110
+ person_mask = cv2.bitwise_not(bg_mask)
111
+
112
+ # Clean up mask with morphological operations
113
+ kernel = np.ones((5, 5), np.uint8)
114
+ person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_CLOSE, kernel)
115
+ person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_OPEN, kernel)
116
+
117
+ # Convert to float and normalize
118
+ return person_mask.astype(np.float32) / 255.0
119
+
120
+ def refine_mask_hq(mask: np.ndarray, frame: np.ndarray, use_matanyone: bool = True) -> np.ndarray:
121
+ """High-quality mask refinement using MatAnyone or fallback methods"""
122
+ try:
123
+ if use_matanyone:
124
+ try:
125
+ global _MATANYONE_PROCESSOR
126
+ if _MATANYONE_PROCESSOR is None:
127
+ _MATANYONE_PROCESSOR = MatAnyoneProcessor()
128
+
129
+ # Ensure proper dtypes
130
+ frame_in = frame if frame.dtype == np.uint8 else frame.astype(np.uint8)
131
+
132
+ # Use MatAnyone to produce a refined alpha matte (0..1 float, HxW)
133
+ alpha = _MATANYONE_PROCESSOR.segment_frame(frame_in, mask_path=None)
134
+
135
+ # Sanity clamp and return
136
+ alpha = np.clip(alpha, 0.0, 1.0).astype(np.float32)
137
+ return alpha
138
+
139
+ except Exception as e:
140
+ logger.warning(f"MatAnyone refinement failed: {e}, using simple refinement")
141
+
142
+ # Fallback: Simple mask refinement
143
+ return _simple_mask_refinement(mask, frame)
144
+
145
+ except Exception as e:
146
+ logger.error(f"Mask refinement failed: {e}")
147
+ return mask
148
+
149
+ def _simple_mask_refinement(mask: np.ndarray, frame: np.ndarray) -> np.ndarray:
150
+ """Simple mask refinement using OpenCV operations"""
151
+ # Convert mask to uint8
152
+ mask_uint8 = (mask * 255).astype(np.uint8)
153
+
154
+ # Apply Gaussian blur for smoother edges
155
+ mask_blurred = cv2.GaussianBlur(mask_uint8, (5, 5), 0)
156
+
157
+ # Apply bilateral filter to preserve edges while smoothing
158
+ mask_refined = cv2.bilateralFilter(mask_blurred, 9, 75, 75)
159
+
160
+ # Convert back to float
161
+ return mask_refined.astype(np.float32) / 255.0
162
+
163
+ def replace_background_hq(frame: np.ndarray, mask: np.ndarray, background: np.ndarray) -> np.ndarray:
164
+ """High-quality background replacement with proper compositing"""
165
+ try:
166
+ # Ensure all inputs are the same size
167
+ h, w = frame.shape[:2]
168
+ background_resized = cv2.resize(background, (w, h))
169
+
170
+ # Ensure mask has 3 channels
171
+ if len(mask.shape) == 2:
172
+ mask_3d = np.stack([mask] * 3, axis=-1)
173
+ else:
174
+ mask_3d = mask
175
+
176
+ # Apply feathering to mask edges for smoother blending
177
+ mask_feathered = _apply_feathering(mask_3d)
178
+
179
+ # Composite the image
180
+ result = frame * mask_feathered + background_resized * (1 - mask_feathered)
181
+
182
+ return result.astype(np.uint8)
183
+
184
+ except Exception as e:
185
+ logger.error(f"Background replacement failed: {e}")
186
+ return frame
187
+
188
+ def _apply_feathering(mask: np.ndarray, feather_amount: int = 3) -> np.ndarray:
189
+ """Apply feathering to mask edges for smoother blending"""
190
+ if len(mask.shape) == 3:
191
+ # Work with single channel
192
+ mask_single = mask[:, :, 0]
193
+ else:
194
+ mask_single = mask
195
+
196
+ # Apply Gaussian blur for feathering
197
+ mask_feathered = cv2.GaussianBlur(mask_single, (feather_amount*2+1, feather_amount*2+1), 0)
198
+
199
+ # Restore 3 channels if needed
200
+ if len(mask.shape) == 3:
201
+ mask_feathered = np.stack([mask_feathered] * 3, axis=-1)
202
+
203
+ return mask_feathered
204
+
205
+ def create_professional_background(bg_type: str, width: int, height: int) -> np.ndarray:
206
+ """Create professional background of specified type and size"""
207
+ try:
208
+ if bg_type not in PROFESSIONAL_BACKGROUNDS:
209
+ bg_type = "office" # Default fallback
210
+
211
+ config = PROFESSIONAL_BACKGROUNDS[bg_type]
212
+ color = config["color"]
213
+ use_gradient = config["gradient"]
214
+
215
+ if use_gradient:
216
+ # Create gradient background
217
+ background = _create_gradient_background(color, width, height)
218
+ else:
219
+ # Create solid color background
220
+ background = np.full((height, width, 3), color, dtype=np.uint8)
221
+
222
+ return background
223
+
224
+ except Exception as e:
225
+ logger.error(f"Background creation failed: {e}")
226
+ # Return white background as fallback
227
+ return np.full((height, width, 3), (255, 255, 255), dtype=np.uint8)
228
+
229
+ def _create_gradient_background(base_color: Tuple[int, int, int], width: int, height: int) -> np.ndarray:
230
+ """Create a gradient background from base color"""
231
+ # Create gradient from darker to lighter version of base color
232
+ r, g, b = base_color
233
+
234
+ # Create darker version (multiply by 0.7)
235
+ dark_color = (int(r * 0.7), int(g * 0.7), int(b * 0.7))
236
+
237
+ # Create gradient
238
+ background = np.zeros((height, width, 3), dtype=np.uint8)
239
+
240
+ for y in range(height):
241
+ # Calculate blend factor (0 to 1)
242
+ blend = y / height
243
+
244
+ # Interpolate between dark and light color
245
+ current_r = int(dark_color[0] * (1 - blend) + r * blend)
246
+ current_g = int(dark_color[1] * (1 - blend) + g * blend)
247
+ current_b = int(dark_color[2] * (1 - blend) + b * blend)
248
+
249
+ background[y, :] = [current_r, current_g, current_b]
250
+
251
+ return background
252
+
253
+ # Export all functions
254
+ __all__ = [
255
+ "segment_person_hq",
256
+ "refine_mask_hq",
257
+ "replace_background_hq",
258
+ "create_professional_background",
259
+ "PROFESSIONAL_BACKGROUNDS",
260
+ "validate_video_file"
261
+ ]