RishiTest / models /misc_utils.py
Lemorra's picture
πŸ”₯ Removed debugging image saving code
6293ac2
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