Upload utils.py

utils.py

@@ -125,56 +125,57 @@ 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 = pillow(Image.open, BytesIO(response.content))
+    img_data = BytesIO(response.content)
 
-    excluded_formats = {'MPO'}
-    output_images, output_masks = [], []
-    w = h = None
+    # 使用 PIL 打开图像
+    img = pillow(Image.open, img_data)
+
+    output_images = []
+    output_masks = []
+    w, h = None, None
     has_alpha = False
 
-    for frame in ImageSequence.Iterator(img):
-        frame = pillow(ImageOps.exif_transpose, frame)
+    excluded_formats = ['MPO']
+
+    for i in ImageSequence.Iterator(img):
+        i = pillow(ImageOps.exif_transpose, i)
 
-        if frame.mode == 'I':
-            frame = frame.point(lambda px: px * (1 / 255))
-        rgb_image = frame.convert("RGB")
+        if i.mode == 'I':
+            i = i.point(lambda i: i * (1 / 255))
+        image = i.convert("RGB")
 
-        if w is None:
-            w, h = rgb_image.size
-        if rgb_image.size != (w, h):
+        if len(output_images) == 0:
+            w, h = image.size
+
+        if image.size != (w, h):
             continue
 
-        image_tensor = torch.from_numpy(np.array(rgb_image, dtype=np.float32) / 255.0)[None, ...]
-        if 'A' in frame.getbands():
+        image_np = np.array(image).astype(np.float32) / 255.0
+        image_tensor = torch.from_numpy(image_np)[None, ...]
+
+        if 'A' in i.getbands():
             has_alpha = True
-            mask_tensor = 1.0 - torch.from_numpy(np.array(frame.getchannel('A'), dtype=np.float32) / 255.0)
+            mask_np = np.array(i.getchannel('A')).astype(np.float32) / 255.0
+            mask_tensor = 1. - torch.from_numpy(mask_np)
         else:
-            mask_tensor = torch.zeros((64, 64), dtype=torch.float32)
+            mask_tensor = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
 
         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_mask = output_images[0], output_masks[0]
+        output_image = output_images[0]
+        output_mask = 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
+        output_mask = output_mask.permute(1, 2, 0).unsqueeze(0)
+        output_image = torch.cat([output_image, output_mask], dim=-1)
 
     return output_image
 
@@ -196,4 +197,9 @@ def extract_json_from_text(text: str) -> dict:
         return json.loads(json_str)
     except json.JSONDecodeError as e:
         print(f"JSON 解析失败: {e}")
-        return {}
+        return {}
+
+
+if __name__ == "__main__":
+    output_image = load_image_from_url("https://obs-large.mtlab.meitu.com/mtopen/fb348748a0ca48cc9ee1ff15059ff499/5264072a-668c-40d5-48c3-624020f88845.png")
+    print(output_image.shape)