返回 VideoClaw
image_processor.py
根目录 / video-claw / video-claw / backend / models / image_processor.py
1 import os
2 import sys
3
4 models_dir = os.path.dirname(os.path.abspath(__file__))
5 backend_dir = os.path.dirname(models_dir)
6 if backend_dir not in sys.path:
7 sys.path.insert(0, backend_dir)
8
9 import requests
10 import numpy as np
11 from pathlib import Path
12 from datetime import datetime, timedelta
13 from PIL import Image
14 import logging
15 from config import Config
16
17 logger = logging.getLogger(__name__)
18
19
20 class ImageProcessor:
21 """
22 图片处理和上传集合类
23 支持:图片处理、分割、拼接,以及上传到阿里云OSS
24 """
25
26 # 阿里云DashScope上传配置
27 UPLOAD_API_URL = "https://dashscope.aliyuncs.com/api/v1/uploads"
28
29 def __init__(self,
30 image_path='',
31 api_key: str = "",
32 model_name: str = "wan2.6-i2v-flash"):
33 """
34 初始化图片处理器
35
36 Args:
37 image_path: 图片文件路径(可选,用于处理已有图片)
38 api_key: DashScope API Key(用于上传,可从 config.yaml 读取)
39 model_name: 模型名称,默认使用 wan2.6-i2v-flash
40 """
41 # 图片处理部分
42 if image_path != '':
43 self.image_path = image_path
44 self.image = Image.open(image_path)
45 self.image_np = np.array(self.image)
46 self.width, self.height = self.image_np.shape[1], self.image_np.shape[0]
47 else:
48 self.image_path = None
49 self.image = None
50 self.image_np = None
51 self.width = None
52 self.height = None
53
54 # 上传功能部分
55 self.api_key = api_key or Config.DASHSCOPE_API_KEY
56 self.model_name = model_name
57
58 @staticmethod
59 def check_column_white(column_pixels):
60 """检查列是否几乎全白"""
61 is_almost_white = np.logical_or(column_pixels == 254, column_pixels == 255)
62 white_pixels_ratio = np.mean(np.all(is_almost_white, axis=-1))
63 return white_pixels_ratio >= 0.98 # 至少98%的像素为白色
64
65 def find_white_section(self, start, end):
66 """查找指定范围内的白色区间"""
67 white_sections = []
68 in_white_section = False
69 start_index = 0
70
71 for col in range(start, end):
72 column_pixels = self.image_np[:, col, :]
73 if self.check_column_white(column_pixels):
74 if not in_white_section:
75 start_index = col
76 in_white_section = True
77 else:
78 if in_white_section:
79 white_sections.append((start_index, col))
80 in_white_section = False
81
82 if in_white_section:
83 white_sections.append((start_index, end))
84
85 return white_sections
86
87 def split_image(self):
88 """将图片从中间分割为左右两部分"""
89 start_col = self.width * 2 // 5
90 end_col = self.width * 3 // 5
91 white_sections = self.find_white_section(start_col, end_col)
92
93 if white_sections:
94 middle_section = white_sections[len(white_sections) // 2]
95 mid_col = (middle_section[0] + middle_section[1]) // 2
96 else:
97 raise ValueError("No suitable white column found within the specified range")
98
99 left_box = (0, 0, mid_col, self.height)
100 right_box = (mid_col, 0, self.width, self.height)
101 left_image = self.image.crop(left_box)
102 right_image = self.image.crop(right_box)
103
104 save_dir, filename = os.path.split(self.image_path)
105 base, extension = os.path.splitext(filename)
106
107 left_image_path = os.path.join(save_dir, base + '_front' + extension)
108 right_image_path = os.path.join(save_dir, base + '_back' + extension)
109 left_image.save(left_image_path)
110 right_image.save(right_image_path)
111
112 return left_image_path, right_image_path
113
114 def stitch_images(self, image_paths, output_path):
115 """拼接多张图片"""
116 if not image_paths:
117 raise ValueError("No image paths provided")
118 sample_image = Image.open(image_paths[0])
119 single_width, single_height = sample_image.size
120 num_images = len(image_paths)
121 total_desired_width = single_width
122 total_current_width = single_width * num_images
123 total_width_to_cut = max(0, total_current_width - total_desired_width)
124 width_to_cut_per_image = total_width_to_cut // num_images
125 stitched_image = Image.new('RGB', (total_desired_width, single_height), "white")
126 current_x = 0
127
128 for path in image_paths:
129 image = Image.open(path)
130 if width_to_cut_per_image > 0:
131 left_margin = width_to_cut_per_image // 2
132 right_margin = image.width - width_to_cut_per_image + left_margin
133 image = image.crop((left_margin, 0, right_margin, image.height))
134 stitched_image.paste(image, (current_x, 0))
135 current_x += image.width
136
137 output_dir = os.path.dirname(output_path)
138 if not os.path.exists(output_dir):
139 os.makedirs(output_dir)
140 stitched_image.save(output_path)
141 return output_path
142
143 def download_image(self, image_url, save_path, max_retries=3, proxies=None):
144 """
145 下载图片,带有重试机制和SSL错误处理
146
147 Args:
148 image_url: 图片URL
149 save_path: 本地保存路径
150 max_retries: 最大重试次数
151 """
152 import time
153 import urllib3
154
155 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
156
157 for attempt in range(max_retries):
158 try:
159 response = requests.get(
160 image_url,
161 timeout=(10, 30),
162 stream=True,
163 verify=True,
164 headers={
165 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
166 },
167 proxies=proxies,
168 )
169
170 if response.status_code == 200:
171 with open(save_path, 'wb') as file:
172 for chunk in response.iter_content(chunk_size=8192):
173 if chunk:
174 file.write(chunk)
175 logger.info("Image downloaded: %s", save_path)
176 return True
177 else:
178 logger.warning("Image download failed: status=%s url=%s", response.status_code, image_url)
179
180 except requests.exceptions.SSLError as e:
181 logger.warning("Image download SSL error: attempt=%d/%d error=%s", attempt + 1, max_retries, str(e)[:100])
182 if attempt < max_retries - 1:
183 wait_time = (attempt + 1) * 2
184 logger.info("Retrying image download in %s seconds", wait_time)
185 time.sleep(wait_time)
186 else:
187 logger.warning("Retrying image download with SSL verification disabled")
188 try:
189 response = requests.get(
190 image_url,
191 timeout=(10, 30),
192 stream=True,
193 verify=False,
194 headers={
195 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
196 },
197 proxies=proxies,
198 )
199 if response.status_code == 200:
200 with open(save_path, 'wb') as file:
201 for chunk in response.iter_content(chunk_size=8192):
202 if chunk:
203 file.write(chunk)
204 logger.info("Image downloaded with SSL verification disabled: %s", save_path)
205 return True
206 except Exception as fallback_error:
207 logger.exception("Image download failed after disabling SSL verification: %s", fallback_error)
208 raise
209
210 except requests.exceptions.Timeout as e:
211 logger.warning("Image download timeout: attempt=%d/%d error=%s", attempt + 1, max_retries, e)
212 if attempt < max_retries - 1:
213 time.sleep((attempt + 1) * 2)
214 else:
215 raise
216
217 except Exception as e:
218 logger.warning("Image download error: attempt=%d/%d error=%s", attempt + 1, max_retries, e)
219 if attempt < max_retries - 1:
220 time.sleep((attempt + 1) * 2)
221 else:
222 raise
223
224 return False
225
226 def resize_image(self, image_path):
227 """调整图片大小(添加顶部空白)"""
228 original_image = Image.open(image_path)
229 width, height = original_image.size
230 top_blank_height = height // 2
231 final_height = height + top_blank_height
232 final_width = int(final_height * 5 / 3)
233 new_image = Image.new("RGB", (final_width, final_height), color="white")
234 left = (final_width - width) // 2
235 top = top_blank_height
236 new_image.paste(original_image, (left, top))
237 new_image.save(image_path)
238 return image_path
239
240 def has_black_borders(self, image_path, threshold=10, black_limit=20):
241 """检查图片是否有黑色边框"""
242 img = Image.open(image_path)
243 pixels = img.load()
244 width, height = img.size
245
246 def is_black_pixel(pixel):
247 return all(x <= black_limit for x in pixel)
248
249 # 检查顶部和底部边框
250 for y in range(threshold):
251 if all(is_black_pixel(pixels[x, y]) for x in range(width)):
252 return True
253 if all(is_black_pixel(pixels[x, height - 1 - y]) for x in range(width)):
254 return True
255
256 # 检查左右边框
257 for x in range(threshold):
258 if all(is_black_pixel(pixels[x, y]) for y in range(height)):
259 return True
260 if all(is_black_pixel(pixels[width - 1 - x, y]) for y in range(height)):
261 return True
262
263 return False
264
265 # ===== 图片上传功能 =====
266
267 def get_upload_policy(self):
268 """
269 获取文件上传凭证
270
271 Returns:
272 policy_data: 包含上传所需凭证的字典
273
274 Raises:
275 Exception: 获取上传凭证失败时
276 """
277 if not self.api_key:
278 raise RuntimeError("DASHSCOPE_API_KEY 未设置,无法使用图片上传服务")
279
280 headers = {
281 "Authorization": f"Bearer {self.api_key}",
282 "Content-Type": "application/json"
283 }
284 params = {
285 "action": "getPolicy",
286 "model": self.model_name
287 }
288
289 response = requests.get(
290 self.UPLOAD_API_URL,
291 headers=headers,
292 params=params,
293 proxies=Config.requests_proxies("dashscope"),
294 )
295 if response.status_code != 200:
296 raise Exception(f"Failed to get upload policy: {response.text}")
297
298 return response.json()['data']
299
300 def upload_file_to_oss(self, policy_data: dict, file_path: str) -> str:
301 """
302 将文件上传到临时存储OSS
303
304 Args:
305 policy_data: 上传凭证数据
306 file_path: 本地文件路径
307
308 Returns:
309 oss_url: OSS URL (格式: oss://...)
310
311 Raises:
312 Exception: 上传失败时
313 """
314 file_name = Path(file_path).name
315 # Sanitize filename for upload to avoid issues with spaces/characters
316 safe_file_name = "".join([c if c.isalnum() or c in ('-','_','.') else '_' for c in file_name])
317
318 key = f"{policy_data['upload_dir']}/{safe_file_name}"
319
320 with open(file_path, 'rb') as file:
321 files = {
322 'OSSAccessKeyId': (None, policy_data['oss_access_key_id']),
323 'Signature': (None, policy_data['signature']),
324 'policy': (None, policy_data['policy']),
325 'x-oss-object-acl': (None, policy_data['x_oss_object_acl']),
326 'x-oss-forbid-overwrite': (None, policy_data['x_oss_forbid_overwrite']),
327 'key': (None, key),
328 'success_action_status': (None, '200'),
329 'file': (safe_file_name, file)
330 }
331
332 response = requests.post(
333 policy_data['upload_host'],
334 files=files,
335 proxies=Config.requests_proxies("dashscope"),
336 )
337 if response.status_code != 200:
338 raise Exception(f"Failed to upload file: {response.text}")
339
340 # Construct OSS URL correctly: oss://<bucket>/<key>
341 # Extract bucket from upload_host (e.g., https://dashscope-instant.oss-cn-beijing.aliyuncs.com)
342 upload_host = policy_data['upload_host']
343 bucket_name = ""
344 if '://' in upload_host:
345 domain = upload_host.split('://')[1]
346 bucket_name = domain.split('.')[0]
347
348 if bucket_name:
349 return f"oss://{bucket_name}/{key}"
350 else:
351 # Fallback if parsing fails (though unlikely for standard OSS hosts)
352 # If the original code's assumption that key was self-sufficient was somehow valid, logic is here.
353 # But normally, oss://<key> is wrong if key doesn't have bucket.
354 return f"oss://{key}"
355
356 def upload(self, file_path: str) -> str:
357 """
358 上传文件到阿里云OSS并获取URL(统一接口方法)
359
360 Args:
361 file_path: 本地文件路径
362
363 Returns:
364 oss_url: OSS URL,可在48小时内使用
365
366 Raises:
367 FileNotFoundError: 文件不存在时
368 RuntimeError: API Key未设置时
369 Exception: 上传失败时
370 """
371 # 检查文件是否存在
372 if not os.path.exists(file_path):
373 raise FileNotFoundError(f"文件不存在: {file_path}")
374
375 if not self.api_key:
376 raise RuntimeError("DASHSCOPE_API_KEY 未设置,无法使用图片上传服务")
377
378 # 1. 获取上传凭证(注意:上传凭证接口有限流)
379 policy_data = self.get_upload_policy()
380
381 # 2. 上传文件到OSS
382 oss_url = self.upload_file_to_oss(policy_data, file_path)
383
384 # 3. 计算过期时间
385 expire_time = datetime.now() + timedelta(hours=48)
386
387 logger.info(
388 "File uploaded to OSS: file=%s oss_url=%s expires_at=%s",
389 file_path,
390 oss_url,
391 expire_time.strftime('%Y-%m-%d %H:%M:%S'),
392 )
393
394 return oss_url
395
396 def collage_images(self, image_paths, output_path):
397 """
398 拼图功能:将多张图片水平拼接
399 Args:
400 image_paths: 图片路径列表
401 output_path: 输出文件路径
402 """
403 if not image_paths:
404 return None
405
406 images = []
407 for p in image_paths:
408 try:
409 img = Image.open(p)
410 images.append(img)
411 except Exception as e:
412 logger.error("Cannot open image %s: %s", p, e)
413
414 if not images:
415 return None
416
417 # 统一高度,按第一张图片的高度调整其他图片
418 base_height = images[0].height
419 resized_images = []
420 for img in images:
421 if img.height != base_height:
422 ratio = base_height / img.height
423 new_width = int(img.width * ratio)
424 resized_images.append(img.resize((new_width, base_height)))
425 else:
426 resized_images.append(img)
427
428 total_width = sum(img.width for img in resized_images)
429 new_im = Image.new('RGB', (total_width, base_height))
430
431 x_offset = 0
432 for img in resized_images:
433 new_im.paste(img, (x_offset, 0))
434 x_offset += img.width
435
436 new_im.save(output_path)
437 return output_path
438
438 lines PYTHON