Spaces:
Sleeping
Sleeping
import cv2 | |
import numpy as np | |
import base64 | |
from PIL import Image | |
from payload_model import PayloadModel | |
def convert_base64_to_cv2(base64_string: str): | |
return cv2.imdecode(np.frombuffer(base64.b64decode(base64_string), np.uint8), cv2.IMREAD_COLOR) | |
def convert_cv2_to_pil(image: np.ndarray): | |
return Image.fromarray(image).convert('RGB') | |
def convert_base64_to_pil(base64_string: str): | |
return convert_cv2_to_pil(convert_base64_to_cv2(base64_string)) | |
def get_images_using_bbox(payload: PayloadModel): | |
images = [] | |
# Forcing that only a single image is received | |
cv2_image = convert_base64_to_cv2(payload.input_data[0]) | |
print(f"Bbox: {payload.bbox}") | |
images_bboxes = payload.bbox | |
image_bboxes = images_bboxes[0] | |
for _, bbox in enumerate(image_bboxes): | |
x1, y1, x2, y2 = bbox | |
image = cv2_image[y1:y2, x1:x2] | |
pil_image = convert_cv2_to_pil(image) | |
images.append(pil_image) | |
return images | |
def get_whole_image(payload: PayloadModel): | |
images = [] | |
# Forcing that only a single image is received | |
pil_image = convert_base64_to_pil(payload.input_data[0]) | |
images.append(pil_image) | |
return images | |