161 lines
5.2 KiB
Python
161 lines
5.2 KiB
Python
import argparse
|
||
import os
|
||
import base64
|
||
import requests
|
||
import json
|
||
from PIL import Image
|
||
from io import BytesIO
|
||
|
||
def load_local_env():
|
||
"""
|
||
Load optional local credentials from config/baidu_api.env.
|
||
This file is ignored by git; see config/baidu_api.env.example.
|
||
"""
|
||
root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
|
||
env_path = os.path.join(root, 'config', 'baidu_api.env')
|
||
if not os.path.exists(env_path):
|
||
return
|
||
with open(env_path, 'r', encoding='utf-8') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line or line.startswith('#') or '=' not in line:
|
||
continue
|
||
key, value = line.split('=', 1)
|
||
os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))
|
||
|
||
|
||
load_local_env()
|
||
|
||
API_KEY = os.environ.get("BAIDU_API_KEY", "")
|
||
SECRET_KEY = os.environ.get("BAIDU_SECRET_KEY", "")
|
||
|
||
# 输入和输出目录
|
||
INPUT_DIR = './去雾图像-北航合作'
|
||
OUTPUT_DIR = './Result_Baidu'
|
||
|
||
def get_access_token():
|
||
"""
|
||
使用 AK，SK 生成鉴权签名（Access Token）
|
||
"""
|
||
if not API_KEY or not SECRET_KEY:
|
||
print("缺少 BAIDU_API_KEY 或 BAIDU_SECRET_KEY，请设置环境变量或 config/baidu_api.env")
|
||
return None
|
||
|
||
url = "https://aip.baidubce.com/oauth/2.0/token"
|
||
params = {
|
||
"grant_type": "client_credentials",
|
||
"client_id": API_KEY,
|
||
"client_secret": SECRET_KEY
|
||
}
|
||
try:
|
||
response = requests.post(url, params=params)
|
||
return response.json().get("access_token")
|
||
except Exception as e:
|
||
print(f"获取 Access Token 失败: {e}")
|
||
return None
|
||
|
||
def process_image(file_path, access_token):
|
||
"""
|
||
读取图片，调用百度API，返回处理后的图片数据
|
||
"""
|
||
request_url = "https://aip.baidubce.com/rest/2.0/image-process/v1/dehaze"
|
||
request_url = request_url + "?access_token=" + access_token
|
||
|
||
try:
|
||
# 使用 PIL 读取图片
|
||
with Image.open(file_path) as img:
|
||
# 将 PIL Image 对象转换为字节流
|
||
img_buffer = BytesIO()
|
||
# 保持原格式保存到内存中
|
||
save_format = img.format if img.format else 'JPEG'
|
||
img.save(img_buffer, format=save_format)
|
||
img_bytes = img_buffer.getvalue()
|
||
|
||
# base64 编码
|
||
img_base64 = base64.b64encode(img_bytes)
|
||
|
||
params = {"image": img_base64}
|
||
headers = {'content-type': 'application/x-www-form-urlencoded'}
|
||
|
||
# 调用 API
|
||
response = requests.post(request_url, data=params, headers=headers)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
# 检查是否有 image 字段
|
||
if "image" in result:
|
||
return base64.b64decode(result["image"])
|
||
else:
|
||
print(f"API 返回错误: {result}")
|
||
return None
|
||
else:
|
||
print(f"请求失败，状态码: {response.status_code}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"处理图片 {file_path} 时出错: {e}")
|
||
return None
|
||
|
||
def main():
|
||
global INPUT_DIR, OUTPUT_DIR
|
||
|
||
parser = argparse.ArgumentParser(description="Run Baidu dehaze API on a folder of images.")
|
||
parser.add_argument("--input-dir", default=INPUT_DIR)
|
||
parser.add_argument("--output-dir", default=OUTPUT_DIR)
|
||
args = parser.parse_args()
|
||
INPUT_DIR = args.input_dir
|
||
OUTPUT_DIR = args.output_dir
|
||
|
||
# 1. 获取 Token
|
||
access_token = get_access_token()
|
||
if not access_token:
|
||
return
|
||
|
||
# 2. 确保输出目录存在
|
||
if not os.path.exists(OUTPUT_DIR):
|
||
os.makedirs(OUTPUT_DIR)
|
||
print(f"已创建输出目录: {OUTPUT_DIR}")
|
||
|
||
# 3. 遍历 ./Data 目录
|
||
if not os.path.exists(INPUT_DIR):
|
||
print(f"输入目录 {INPUT_DIR} 不存在")
|
||
return
|
||
|
||
supported_exts = ('.png', '.bmp', '.jpg', '.jpeg')
|
||
|
||
print("开始批量处理...")
|
||
|
||
files = os.listdir(INPUT_DIR)
|
||
for filename in files:
|
||
# 检查文件扩展名
|
||
if filename.lower().endswith(supported_exts):
|
||
input_path = os.path.join(INPUT_DIR, filename)
|
||
# V1
|
||
# output_path = os.path.join(OUTPUT_DIR, filename)
|
||
# V2
|
||
# 1. 分离文件名和扩展名
|
||
name, ext = os.path.splitext(filename)
|
||
# 2. 拼接新文件名
|
||
new_filename = f"{name}_result{ext}"
|
||
# 3. 生成最终输出路径
|
||
output_path = os.path.join(OUTPUT_DIR, new_filename)
|
||
|
||
|
||
print(f"正在处理: {filename} ...")
|
||
|
||
# 处理图片
|
||
processed_data = process_image(input_path, access_token)
|
||
|
||
# 保存结果
|
||
if processed_data:
|
||
with open(output_path, 'wb') as f:
|
||
f.write(processed_data)
|
||
print(f" 已保存至: {output_path}")
|
||
else:
|
||
print(f" 处理失败: {filename}")
|
||
|
||
print("批量处理完成。")
|
||
|
||
if __name__ == '__main__':
|
||
main()
|