Spaces:
Sleeping
Sleeping
| from PIL import Image | |
| def predict(yolo_result : any | |
| ,img : Image = None) -> int: | |
| """ | |
| pedicts the orientation based on the position of the largest table anotation and returns a degree the image should be rotated. | |
| input: | |
| yolo_result : for one image | |
| img : PIL Image, default None | |
| if None the image from result is taken | |
| """ | |
| result = yolo_result | |
| boxes = result.boxes # Boxes object for bounding box outputs | |
| # masks = result.masks # Masks object for segmentation masks outputs | |
| # keypoints = result.keypoints # Keypoints object for pose outputs | |
| # probs = result.probs # Probs object for classification outputs | |
| if img is None: | |
| img = result.plot() # BGR-order numpy array | |
| img = Image.fromarray(img[..., ::-1]) # RGB-order PIL image | |
| box_with_max_volume = get_reference_table(boxes) | |
| if box_with_max_volume == None: | |
| return 0 #early return if nothing found. | |
| # Get the coordinates of the box with the largest volume | |
| x1, y1, x2, y2 = box_with_max_volume.xyxy[0].tolist() | |
| # Get the distances to the borders | |
| dist_top = y1 | |
| dist_left = x1 | |
| dist_right = img.width - x2 | |
| dist_bottom = img.height - y2 | |
| # Determine the rotation angle based on the distances to the borders | |
| if dist_top < dist_bottom and dist_left < dist_right: # top left corner | |
| rotation_angle = 180 # Rotate by 180 degrees | |
| elif dist_top < dist_bottom and dist_left >= dist_right: # top right corner | |
| rotation_angle = -90 # Rotate by -90 degrees | |
| elif dist_top > dist_bottom and dist_left < dist_right: # bottom left | |
| rotation_angle = 90 | |
| else: | |
| rotation_angle = 0 # do nothing | |
| return rotation_angle # , box_with_max_volume | |
| # Rotate the original image by the calculated angle | |
| # rotated_image = img.rotate(rotation_angle, expand=True) | |
| def get_reference_table(boxes): | |
| """ | |
| Returns the reference table from result | |
| """ | |
| # get all tables by class id | |
| tables = [box for box in boxes if box.cls.item() == 1] | |
| if len(tables) == 0: | |
| return None | |
| # detect box with the largest volume | |
| box_with_max_volume = max( | |
| tables, | |
| key=lambda box: (box.xywhn[0][2].item() * box.xywhn[0][3].item()) | |
| ) | |
| # retun the biggest box | |
| return box_with_max_volume |