Spaces:
Paused
Paused
| 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 | |