Update uploader/do_spaces.py

#2
Files changed (1) hide show
  1. uploader/do_spaces.py +14 -14
uploader/do_spaces.py CHANGED
@@ -7,27 +7,27 @@ import io
7
 
8
  def upload_mask(image, prefix="mask"):
9
  """
10
- 将分割掩码图像上传到DigitalOcean Spaces
11
 
12
  Args:
13
- image: PIL Image对象
14
- prefix: 文件名前缀
15
 
16
  Returns:
17
- 上传文件的URL
18
  """
19
  try:
20
- # 从环境变量获取凭据
21
  do_key = os.environ.get('DO_SPACES_KEY')
22
  do_secret = os.environ.get('DO_SPACES_SECRET')
23
  do_region = os.environ.get('DO_SPACES_REGION')
24
  do_bucket = os.environ.get('DO_SPACES_BUCKET')
25
 
26
- # 校验凭据是否存在
27
  if not all([do_key, do_secret, do_region, do_bucket]):
28
- raise ValueError("缺少DigitalOcean Spaces凭据")
29
 
30
- # 创建S3客户端
31
  session = boto3.session.Session()
32
  client = session.client('s3',
33
  region_name=do_region,
@@ -35,15 +35,15 @@ def upload_mask(image, prefix="mask"):
35
  aws_access_key_id=do_key,
36
  aws_secret_access_key=do_secret)
37
 
38
- # 生成唯一文件名
39
  filename = f"{prefix}_{uuid.uuid4().hex}.png"
40
 
41
- # 将图像转换为字节流
42
  img_byte_arr = io.BytesIO()
43
  image.save(img_byte_arr, format='PNG')
44
  img_byte_arr.seek(0)
45
 
46
- # 上传到Spaces
47
  client.upload_fileobj(
48
  img_byte_arr,
49
  do_bucket,
@@ -51,10 +51,10 @@ def upload_mask(image, prefix="mask"):
51
  ExtraArgs={'ACL': 'public-read', 'ContentType': 'image/png'}
52
  )
53
 
54
- # 返回公共URL
55
  url = f'https://{do_bucket}.{do_region}.digitaloceanspaces.com/{filename}'
56
  return url
57
 
58
  except Exception as e:
59
- print(f"上传失败: {str(e)}")
60
- return f"上传错误: {str(e)}"
 
7
 
8
  def upload_mask(image, prefix="mask"):
9
  """
10
+ Upload the segmentation mask image to DigitalOcean Spaces
11
 
12
  Args:
13
+ image: PIL Image object
14
+ prefix: filename prefix
15
 
16
  Returns:
17
+ URL of the uploaded file
18
  """
19
  try:
20
+ # Get credentials from environment variables
21
  do_key = os.environ.get('DO_SPACES_KEY')
22
  do_secret = os.environ.get('DO_SPACES_SECRET')
23
  do_region = os.environ.get('DO_SPACES_REGION')
24
  do_bucket = os.environ.get('DO_SPACES_BUCKET')
25
 
26
+ # Validate that credentials exist
27
  if not all([do_key, do_secret, do_region, do_bucket]):
28
+ raise ValueError("Missing DigitalOcean Spaces credentials")
29
 
30
+ # Create S3 client
31
  session = boto3.session.Session()
32
  client = session.client('s3',
33
  region_name=do_region,
 
35
  aws_access_key_id=do_key,
36
  aws_secret_access_key=do_secret)
37
 
38
+ # Generate a unique filename
39
  filename = f"{prefix}_{uuid.uuid4().hex}.png"
40
 
41
+ # Convert the image to a byte stream
42
  img_byte_arr = io.BytesIO()
43
  image.save(img_byte_arr, format='PNG')
44
  img_byte_arr.seek(0)
45
 
46
+ # Upload to Spaces
47
  client.upload_fileobj(
48
  img_byte_arr,
49
  do_bucket,
 
51
  ExtraArgs={'ACL': 'public-read', 'ContentType': 'image/png'}
52
  )
53
 
54
+ # Return the public URL
55
  url = f'https://{do_bucket}.{do_region}.digitaloceanspaces.com/{filename}'
56
  return url
57
 
58
  except Exception as e:
59
+ print(f"Upload failed: {str(e)}")
60
+ return f"Upload error: {str(e)}"