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, )