Upload utils.py

utils.py

@@ -125,51 +125,56 @@ def tensor_to_base64(img_tensor: torch.Tensor) -> str:
 
 
 def load_image_from_url(image_url):
-    # 下载图片数据
+    """从 URL 加载图像为 torch.Tensor，支持多帧和 Alpha 通道融合。"""
+    # 下载并读取图像
     response = requests.get(image_url)
     response.raise_for_status()
-    img_data = BytesIO(response.content)
+    img = pillow(Image.open, BytesIO(response.content))
 
-    # 使用 PIL 打开图像
-    img = pillow(Image.open, img_data)
+    excluded_formats = {'MPO'}
+    output_images, output_masks = [], []
+    w = h = None
+    has_alpha = False
 
-    output_images = []
-    output_masks = []
-    w, h = None, None
+    for frame in ImageSequence.Iterator(img):
+        frame = pillow(ImageOps.exif_transpose, frame)
 
-    excluded_formats = ['MPO']
+        if frame.mode == 'I':
+            frame = frame.point(lambda px: px * (1 / 255))
+        rgb_image = frame.convert("RGB")
 
-    for i in ImageSequence.Iterator(img):
-        i = pillow(ImageOps.exif_transpose, i)
-
-        if i.mode == 'I':
-            i = i.point(lambda i: i * (1 / 255))
-        image = i.convert("RGB")
-
-        if len(output_images) == 0:
-            w, h = image.size
-
-        if image.size != (w, h):
+        if w is None:
+            w, h = rgb_image.size
+        if rgb_image.size != (w, h):
             continue
 
-        image_np = np.array(image).astype(np.float32) / 255.0
-        image_tensor = torch.from_numpy(image_np)[None, ...]
-
-        if 'A' in i.getbands():
-            mask_np = np.array(i.getchannel('A')).astype(np.float32) / 255.0
-            mask_tensor = 1. - torch.from_numpy(mask_np)
+        image_tensor = torch.from_numpy(np.array(rgb_image, dtype=np.float32) / 255.0)[None, ...]
+        if 'A' in frame.getbands():
+            has_alpha = True
+            mask_tensor = 1.0 - torch.from_numpy(np.array(frame.getchannel('A'), dtype=np.float32) / 255.0)
         else:
-            mask_tensor = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
+            mask_tensor = torch.zeros((64, 64), dtype=torch.float32)
 
         output_images.append(image_tensor)
         output_masks.append(mask_tensor.unsqueeze(0))
 
+    # 合并帧
     if len(output_images) > 1 and getattr(img, "format", None) not in excluded_formats:
         output_image = torch.cat(output_images, dim=0)
         output_mask = torch.cat(output_masks, dim=0)
     else:
-        output_image = output_images[0]
-        output_mask = output_masks[0]
+        output_image, output_mask = output_images[0], output_masks[0]
+
+    # 融合 Alpha 通道
+    if has_alpha:
+        image = output_image.squeeze(0)
+        if image.dim() == 3 and image.shape[0] in (1, 3, 4):
+            image = image.permute(1, 2, 0)
+        h, w = output_mask.shape
+        rgba = torch.zeros(h, w, 4)
+        rgba[:, :, :3] = image
+        rgba[:, :, 3] = output_mask
+        output_image = rgba
 
     return output_image