from PIL import Image
import os
import time
import requests
import base64
import numpy as np
from io import BytesIO
import concurrent.futures

NUM_SEGMENTS = 10

tic = time.time()
video_dir = '/home/ubuntu/shared_storage/images'
frames = [(os.path.splitext(item)[0], os.path.join(video_dir, item)) for item in os.listdir(video_dir)]
frames = [item[1] for item in sorted(frames, key=lambda x: x[0])]
indices = np.linspace(start=0, stop=len(frames)-1, num=NUM_SEGMENTS).astype(int)
image_paths = [frames[ind] for ind in indices]

request = {}
byte_images = []

def image_path_handler(image_path):
    img = Image.open(image_path)
    byte_io = BytesIO()
    img.save(byte_io, format='PNG')
    encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
    return encoded_image

with concurrent.futures.ThreadPoolExecutor() as executor:
    byte_images = list(executor.map(image_path_handler, image_paths))

payload = {
        "images": byte_images * 16,
        "parameters": {
            "max_new_tokens": 90,
            "top_k": 4,
            "top_p": None,
            "temperature": 0.01,
            "no_repeat_ngram_size": None,
        }
    }

IP = '127.0.0.1'
url = f'http://{IP}:8000/generate'
headers = {'Content-Type': 'application/json'}

tic1 = time.time()

response = requests.post(url, json=payload, headers=headers)
toc = time.time()
print(response.text)

print('stage 1', tic1 - tic)
print('stage 2', toc - tic1)