Background Removal

1
import requests
2
import json
3
import os
4
import base64
5
import time
6
import io
7
import PIL.Image
8
from typing import Optional, Tuple
9
10
11
def _process_test(url):
12
    
13
    image = PIL.Image.open("headphones1.jpeg")
14
15
    # Create a BytesIO buffer to hold the image data
16
    image_buffer = io.BytesIO()
17
    image.save(image_buffer, format='JPEG')
18
    image_bytes = image_buffer.getvalue()
19
    encoded_image = base64.b64encode(image_bytes).decode('utf-8')
20
21
    OCTOAI_TOKEN = os.environ.get("OCTOAI_TOKEN")
22
23
    payload = {
24
        "init_image": encoded_image,
25
        "bgcolor":(255, 255, 255, 0)
26
    }
27
    headers = {
28
        "Authorization": f"Bearer {OCTOAI_TOKEN}",
29
        "Content-Type": "application/json",
30
        "X-OctoAI-Queue-Dispatch": "true"
31
    }
32
33
    response = requests.post(url, headers=headers, json=payload)
34
35
    if response.status_code != 200:
36
        print(response.text)
37
    print(response.json())
38
39
    img_info = response.json()["image_b64"]
40
41
    img_bytes = base64.b64decode(img_info)
42
    img = PIL.Image.open(io.BytesIO(img_bytes))
43
44
    if img.mode == 'RGBA':
45
     img = img.convert('RGB')
46
47
    img.save("result_image.png")
48
49
if __name__ == "__main__":
50
    _process_test("https://image.octoai.run/background-removal")