import io
import time
import zipfile
import requests
from PIL import Image, ImageColor
from carvekit.utils.image_utils import transparency_paste, add_margin
from carvekit.utils.mask_utils import extract_alpha_channel
from carvekit.web.responses.api import error_dict
from carvekit.api.interface import Interface
def process_remove_bg(
interface: Interface, params, image, bg, is_json_or_www_encoded=False
):
"""
Handles a request to the removebg api method
Args:
interface: CarveKit interface
bg: background pil image
is_json_or_www_encoded: is "json" or "x-www-form-urlencoded" content-type
image: foreground pil image
params: parameters
"""
h, w = image.size
if h < 2 or w < 2:
return error_dict("Image is too small. Minimum size 2x2"), 400
if "size" in params.keys():
value = params["size"]
if value == "preview" or value == "small" or value == "regular":
image.thumbnail((625, 400), resample=3) # 0.25 mp
elif value == "medium":
image.thumbnail((1504, 1000), resample=3) # 1.5 mp
elif value == "hd":
image.thumbnail((2000, 2000), resample=3) # 2.5 mp
else:
image.thumbnail((6250, 4000), resample=3) # 25 mp
roi_box = [0, 0, image.size[0], image.size[1]]
if "type" in params.keys():
value = params["type"]
pass
if "roi" in params.keys():
value = params["roi"].split(" ")
if len(value) == 4:
for i, coord in enumerate(value):
if "px" in coord:
coord = coord.replace("px", "")
try:
coord = int(coord)
except BaseException:
return (
error_dict(
"Error converting roi coordinate string to number!"
),
400,
)
if coord < 0:
error_dict("Bad roi coordinate."), 400
if (i == 0 or i == 2) and coord > image.size[0]:
return (
error_dict(
"The roi coordinate cannot be larger than the image size."
),
400,
)
elif (i == 1 or i == 3) and coord > image.size[1]:
return (
error_dict(
"The roi coordinate cannot be larger than the image size."
),
400,
)
roi_box[i] = int(coord)
elif "%" in coord:
coord = coord.replace("%", "")
try:
coord = int(coord)
except BaseException:
return (
error_dict(
"Error converting roi coordinate string to number!"
),
400,
)
if coord > 100:
return (
error_dict("The coordinate cannot be more than 100%"),
400,
)
elif coord < 0:
return error_dict("Coordinate cannot be less than 0%"), 400
if i == 0 or i == 2:
coord = int(image.size[0] * coord / 100)
elif i == 1 or i == 3:
coord = int(image.size[1] * coord / 100)
roi_box[i] = coord
else:
return error_dict("Something wrong with roi coordinates!"), 400
new_image = image.copy()
new_image = new_image.crop(roi_box)
h, w = new_image.size
if h < 2 or w < 2:
return error_dict("Image is too small. Minimum size 2x2"), 400
new_image = interface([new_image])[0]
scaled = False
if "scale" in params.keys() and params["scale"] != 100:
value = params["scale"]
new_image.thumbnail(
(int(image.size[0] * value / 100), int(image.size[1] * value / 100)),
resample=3,
)
scaled = True
if "crop" in params.keys():
value = params["crop"]
if value:
new_image = new_image.crop(new_image.getbbox())
if "crop_margin" in params.keys():
crop_margin = params["crop_margin"]
if "px" in crop_margin:
crop_margin = crop_margin.replace("px", "")
crop_margin = abs(int(crop_margin))
if crop_margin > 500:
return (
error_dict(
"The crop_margin cannot be larger than the original image size."
),
400,
)
new_image = add_margin(
new_image,
crop_margin,
crop_margin,
crop_margin,
crop_margin,
(0, 0, 0, 0),
)
elif "%" in crop_margin:
crop_margin = crop_margin.replace("%", "")
crop_margin = int(crop_margin)
new_image = add_margin(
new_image,
int(new_image.size[1] * crop_margin / 100),
int(new_image.size[0] * crop_margin / 100),
int(new_image.size[1] * crop_margin / 100),
int(new_image.size[0] * crop_margin / 100),
(0, 0, 0, 0),
)
else:
if "position" in params.keys() and scaled is False:
value = params["position"]
if len(value) == 2:
new_image = transparency_paste(
Image.new("RGBA", image.size),
new_image,
(
int(image.size[0] * value[0] / 100),
int(image.size[1] * value[1] / 100),
),
)
else:
new_image = transparency_paste(
Image.new("RGBA", image.size), new_image, roi_box
)
elif scaled is False:
new_image = transparency_paste(
Image.new("RGBA", image.size), new_image, roi_box
)
if "channels" in params.keys():
value = params["channels"]
if value == "alpha":
new_image = extract_alpha_channel(new_image)
else:
bg_changed = False
if "bg_color" in params.keys():
value = params["bg_color"]
if len(value) > 0:
color = ImageColor.getcolor(value, "RGB")
bg = Image.new("RGBA", new_image.size, color)
bg = transparency_paste(bg, new_image, (0, 0))
new_image = bg.copy()
bg_changed = True
if "bg_image_url" in params.keys() and bg_changed is False:
value = params["bg_image_url"]
if len(value) > 0:
try:
bg = Image.open(io.BytesIO(requests.get(value).content))
except BaseException:
return error_dict("Error download background image!"), 400
bg = bg.resize(new_image.size)
bg = bg.convert("RGBA")
bg = transparency_paste(bg, new_image, (0, 0))
new_image = bg.copy()
bg_changed = True
if not is_json_or_www_encoded:
if bg and bg_changed is False:
bg = bg.resize(new_image.size)
bg = bg.convert("RGBA")
bg = transparency_paste(bg, new_image, (0, 0))
new_image = bg.copy()
if "format" in params.keys():
value = params["format"]
if value == "jpg":
new_image = new_image.convert("RGB")
img_io = io.BytesIO()
new_image.save(img_io, "JPEG", quality=100)
img_io.seek(0)
return {"type": "jpg", "data": [img_io, new_image.size]}
elif value == "zip":
mask = extract_alpha_channel(new_image)
mask_buff = io.BytesIO()
mask.save(mask_buff, "PNG")
mask_buff.seek(0)
image_buff = io.BytesIO()
image.save(image_buff, "JPEG")
image_buff.seek(0)
fileobj = io.BytesIO()
with zipfile.ZipFile(fileobj, "w") as zip_file:
zip_info = zipfile.ZipInfo(filename="color.jpg")
zip_info.date_time = time.localtime(time.time())[:6]
zip_info.compress_type = zipfile.ZIP_DEFLATED
zip_file.writestr(zip_info, image_buff.getvalue())
zip_info = zipfile.ZipInfo(filename="alpha.png")
zip_info.date_time = time.localtime(time.time())[:6]
zip_info.compress_type = zipfile.ZIP_DEFLATED
zip_file.writestr(zip_info, mask_buff.getvalue())
fileobj.seek(0)
return {"type": "zip", "data": [fileobj.read(), new_image.size]}
else:
buff = io.BytesIO()
new_image.save(buff, "PNG")
buff.seek(0)
return {"type": "png", "data": [buff, new_image.size]}
return (
error_dict(
"Something wrong with request or http api. Please, open new issue on Github! This is error in "
"code."
),
400,
)