64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
from flask import Flask, request
|
|
from PIL import Image, ImageOps
|
|
import requests
|
|
import io
|
|
import os
|
|
import base64
|
|
from flask_cors import CORS
|
|
|
|
app = Flask(__name__)
|
|
CORS(app)
|
|
|
|
@app.route('/get_item_mask', methods=['POST'])
|
|
def get_item_mask():
|
|
|
|
# use remove bg to get a png image and convert it to a black and white mask image.
|
|
print("get item mask called")
|
|
|
|
image_file = request.files['image']
|
|
image_bytes = image_file.read()
|
|
image = Image.open(io.BytesIO(image_bytes))
|
|
|
|
image = image.convert('RGBA')
|
|
image.save('input.png')
|
|
|
|
jpeg_buffer = io.BytesIO()
|
|
|
|
ai_mask_image = convert_to_black_and_white(image)
|
|
|
|
ai_mask_image.save("output.png")
|
|
ai_mask_image = ai_mask_image.convert('RGB')
|
|
ai_mask_image.save(jpeg_buffer, format="JPEG")
|
|
ai_mask_image_base64 = base64.b64encode(jpeg_buffer.getvalue()).decode('utf-8')
|
|
|
|
return {'ai_mask': ai_mask_image_base64}
|
|
|
|
|
|
def convert_to_black_and_white(image):
|
|
# Create a new image with the same size as the original, but with a white background
|
|
new_image = Image.new('RGB', image.size, (255, 255, 255))
|
|
|
|
# Ensure the image has an alpha channel
|
|
if image.mode != 'RGBA':
|
|
image = image.convert('RGBA')
|
|
|
|
# Get pixel data
|
|
data = image.getdata()
|
|
|
|
# Create new image data
|
|
new_data = []
|
|
for item in data:
|
|
# If the pixel is fully transparent (alpha value of 0), set it to white. Else, set it to black.
|
|
if item[3] == 0:
|
|
new_data.append((255, 255, 255, 255))
|
|
else:
|
|
new_data.append((0, 0, 0, 255))
|
|
|
|
# Update new image data
|
|
new_image.putdata(new_data)
|
|
|
|
# Convert the final image to grayscale for black and white effect
|
|
black_and_white_image = new_image.convert('L')
|
|
|
|
return black_and_white_image
|