|
|
from PIL import Image |
|
|
import torch |
|
|
from torchvision import transforms |
|
|
from transformers import AutoModelForImageSegmentation |
|
|
from typing import Dict, List, Any |
|
|
import base64 |
|
|
from io import BytesIO |
|
|
import os |
|
|
import boto3 |
|
|
import datetime |
|
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
class EndpointHandler(): |
|
|
def __init__(self, path=""): |
|
|
self.model = AutoModelForImageSegmentation.from_pretrained( |
|
|
'whlzy/remove_bg_api', |
|
|
trust_remote_code=True, |
|
|
token=os.environ.get("HUGGINGFACE_TOKEN") |
|
|
) |
|
|
self.model.to(device) |
|
|
self.model.eval() |
|
|
image_size = (1024, 1024) |
|
|
self.transform_image = transforms.Compose([ |
|
|
transforms.Resize(image_size), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
|
|
]) |
|
|
|
|
|
def update_to_s3(self, image): |
|
|
BUCKET_NAME = 'popwear-assets' |
|
|
BUCKET_PREFIX_PATH = 'removebg' |
|
|
ACCOUNT_ID = '18cc2282d0ee72171c1ea322ed22983c' |
|
|
ACCESS_KEY_ID = '007f1852a377a2df43a21d5c8d54542e' |
|
|
SECRET_ACCESS_KEY = 'db2658e2429950bb05e15afb6c53c8b7fd23ab9e1bf79cd42604c89f276068e4' |
|
|
ENDPOINT_URL = f'https://{ACCOUNT_ID}.r2.cloudflarestorage.com' |
|
|
bucket_postfix_path = f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg" |
|
|
image_url = f"https://assets.popwear.ai/{BUCKET_PREFIX_PATH}/{bucket_postfix_path}" |
|
|
s3 = boto3.client( |
|
|
's3', |
|
|
endpoint_url=ENDPOINT_URL, |
|
|
aws_access_key_id=ACCESS_KEY_ID, |
|
|
aws_secret_access_key=SECRET_ACCESS_KEY, |
|
|
region_name='auto' |
|
|
) |
|
|
output_buffer = BytesIO() |
|
|
image.save(output_buffer, format='WEBP', quality=85, method=4) |
|
|
output_buffer.seek(0) |
|
|
s3.upload_fileobj(output_buffer, BUCKET_NAME, f"{BUCKET_PREFIX_PATH}/{bucket_postfix_path}") |
|
|
return image_url |
|
|
|
|
|
def __call__(self, data: Any) -> List[List[Dict[str, float]]]: |
|
|
image = data.pop("inputs", data) |
|
|
|
|
|
input_images = self.transform_image(image).unsqueeze(0).to('cuda') |
|
|
with torch.no_grad(): |
|
|
preds = self.model(input_images)[-1].sigmoid().cpu() |
|
|
pred = preds[0].squeeze() |
|
|
pred_pil = transforms.ToPILImage()(pred) |
|
|
mask = pred_pil.resize(image.size) |
|
|
image.putalpha(mask) |
|
|
image_url = self.update_to_s3(image) |
|
|
return image_url |
|
|
|
|
|
def decode_base64_image(self, image_string): |
|
|
base64_image = base64.b64decode(image_string) |
|
|
buffer = BytesIO(base64_image) |
|
|
image = Image.open(buffer) |
|
|
return image |
|
|
|
