File size: 3,915 Bytes
1e3b872 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import aiohttp
from PIL import Image
import io
import torch
import numpy as np
import base64
import json
import asyncio
import ssl
ssl_context = ssl.create_default_context()
class HumanParseNode:
CATEGORY = "WEFA-DOOR"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"input_image": ("IMAGE", {}),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
RETURN_NAMES = ("top_mask", "bottom_mask", "combined_map")
FUNCTION = "send_to_api_sync"
@staticmethod
async def image_to_base64(image_tensor):
# Convert PyTorch tensor to numpy array
image_np = image_tensor.cpu().detach().numpy()
# Handle the case for a single image in a batch
if image_np.ndim == 4 and image_np.shape[0] == 1:
# Remove the batch dimension if it's a single image in a batch
image_np = image_np.squeeze(0)
# If the image is in CHW format, convert it to HWC
if image_np.ndim == 3 and image_np.shape[0] in {1, 3}:
# Assume CHW format and convert to HWC
image_np = image_np.transpose(1, 2, 0)
# Ensure the image data is in uint8
image_np = (image_np * 255).astype(np.uint8)
# Handle the case where we have a grayscale image (H, W, 1) -> (H, W)
if image_np.ndim == 3 and image_np.shape[2] == 1:
image_np = image_np.squeeze(2)
# Create PIL Image and encode to base64
image_pil = Image.fromarray(image_np)
buffer = io.BytesIO()
image_pil.save(buffer, format="PNG")
return base64.b64encode(buffer.getvalue()).decode()
async def send_to_api(self, input_image):
url = 'http://63.141.33.9:22161'
headers = {
'X-API-KEY': 'xiCQTaoQKXUNATzuFLWRgtoJKiFXiDGvnk',
'Content-Type': 'application/json'
}
json_payload = json.dumps({
"image": await HumanParseNode.image_to_base64(input_image)
})
# print("Sending JSON payload:", json_payload)
# Ensure aiohttp uses the SSL context, particularly important for handling HTTPS requests securely
async with aiohttp.ClientSession() as session:
async with session.post(url, data=json_payload, headers=headers, ssl=ssl_context) as response:
print(f"Response status: {response.status}")
if response.status != 200:
response_text = await response.text()
print(f"Error: Response status {response.status}. Response text: {response_text}")
return None
else:
response_json = await response.json()
top_mask_base64 = response_json.get("1.png")
bottom_mask_base64 = response_json.get("2.png")
combined_map_base64 = response_json.get("final_seg.png")
top_mask_image = Image.open(io.BytesIO(base64.b64decode(top_mask_base64))).convert('L')
bottom_mask_image = Image.open(io.BytesIO(base64.b64decode(bottom_mask_base64))).convert('L')
combined_map_image = Image.open(io.BytesIO(base64.b64decode(combined_map_base64))).convert('RGB')
top_mask_tensor = torch.tensor(np.array(top_mask_image)).unsqueeze(0).float() / 255.0
bottom_mask_tensor = torch.tensor(np.array(bottom_mask_image)).unsqueeze(0).float() / 255.0
combined_map_tensor = torch.tensor(np.array(combined_map_image)).unsqueeze(0).float() / 255.0
return (top_mask_tensor, bottom_mask_tensor, combined_map_tensor)
def send_to_api_sync(self, input_image):
# This wrapper is synchronous and will block until the coroutine has finished executing.
return asyncio.run(self.send_to_api(input_image))
|