|
|
from pathlib import Path |
|
|
from typing import Optional, Union |
|
|
import numpy as np |
|
|
from ultralytics import YOLO |
|
|
from ultralyticsplus import YOLO as YOLO2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TableDetector: |
|
|
""" |
|
|
A class for detecting tables in document images using YOLO models. |
|
|
|
|
|
Attributes: |
|
|
model_path (Path): Path to the YOLO model weights |
|
|
confidence (float): Confidence threshold for detection |
|
|
iou_threshold (float): IoU threshold for NMS |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
confidence: float = 0.50, |
|
|
iou_threshold: float = 0.45 |
|
|
) -> None: |
|
|
""" |
|
|
Initialize the TableDetector with model and parameters. |
|
|
|
|
|
Args: |
|
|
model_path: Path to the YOLO model weights |
|
|
confidence: Confidence threshold for detection |
|
|
iou_threshold: IoU threshold for NMS |
|
|
""" |
|
|
self.model_path = 'src/models/table-detection-and-extraction.pt' |
|
|
self.model = YOLO(str(self.model_path)) |
|
|
self.min_conf = confidence |
|
|
self.iou = iou_threshold |
|
|
|
|
|
def load_model(self): |
|
|
model = YOLO2('foduucom/table-detection-and-extraction') |
|
|
model.overrides['conf'] = 0.25 |
|
|
model.overrides['iou'] = 0.45 |
|
|
model.overrides['agnostic_nms'] = False |
|
|
model.overrides['max_det'] = 1000 |
|
|
return model |
|
|
|
|
|
def detect(self, image_path: Union[str, Path]) -> Optional[np.ndarray]: |
|
|
""" |
|
|
Detect tables in the given image. |
|
|
|
|
|
Args: |
|
|
image_path: Path to the input image |
|
|
|
|
|
Returns: |
|
|
Array of bounding box coordinates or None if no tables detected |
|
|
""" |
|
|
results = self.model.predict(str(image_path), verbose=False, iou = self.iou, conf = self.min_conf) |
|
|
if results: |
|
|
boxes = results[0].boxes.xyxy.numpy() |
|
|
cord = self.merge_boxes(boxes) |
|
|
return [sorted(cord, key = lambda x : (x[2]-x[0])* (x[3]-x[1]), reverse=True)[0]] if len(cord) > 0 else [] |
|
|
|
|
|
return None |
|
|
|
|
|
def merge_boxes(self, boxes: np.ndarray, overlap_threshold: float = 35) -> np.ndarray: |
|
|
""" |
|
|
Merge overlapping bounding boxes. |
|
|
|
|
|
Args: |
|
|
boxes: Array of bounding box coordinates |
|
|
overlap_threshold: Threshold for merging overlapping boxes |
|
|
|
|
|
Returns: |
|
|
Array of merged bounding box coordinates |
|
|
""" |
|
|
|
|
|
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) |
|
|
sorted_indices = np.argsort(-areas) |
|
|
boxes = boxes[sorted_indices] |
|
|
|
|
|
merged_boxes = [] |
|
|
|
|
|
for box in boxes: |
|
|
if not merged_boxes: |
|
|
merged_boxes.append(box) |
|
|
continue |
|
|
|
|
|
overlap_found = False |
|
|
for i, merged_box in enumerate(merged_boxes): |
|
|
iou = self._calculate_overlap(box, merged_box) |
|
|
if iou > overlap_threshold: |
|
|
|
|
|
box_area = (box[2] - box[0]) * (box[3] - box[1]) |
|
|
merged_area = (merged_box[2] - merged_box[0]) * (merged_box[3] - merged_box[1]) |
|
|
if box_area > merged_area: |
|
|
merged_boxes[i] = box |
|
|
overlap_found = True |
|
|
break |
|
|
|
|
|
if not overlap_found: |
|
|
merged_boxes.append(box) |
|
|
|
|
|
return np.array(merged_boxes).astype(int) |
|
|
|
|
|
@staticmethod |
|
|
def _calculate_overlap(box1: np.ndarray, box2: np.ndarray) -> float: |
|
|
""" |
|
|
Calculate the percentage overlap between two boxes. |
|
|
|
|
|
Args: |
|
|
box1: First bounding box coordinates |
|
|
box2: Second bounding box coordinates |
|
|
|
|
|
Returns: |
|
|
Percentage of overlap between the boxes |
|
|
""" |
|
|
x_left = max(box1[0], box2[0]) |
|
|
y_top = max(box1[1], box2[1]) |
|
|
x_right = min(box1[2], box2[2]) |
|
|
y_bottom = min(box1[3], box2[3]) |
|
|
|
|
|
if x_right < x_left or y_bottom < y_top: |
|
|
return 0.0 |
|
|
|
|
|
intersection_area = (x_right - x_left) * (y_bottom - y_top) |
|
|
box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) |
|
|
box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) |
|
|
|
|
|
min_area = min(box1_area, box2_area) |
|
|
if min_area == 0: |
|
|
return 0.0 |
|
|
|
|
|
return (intersection_area / min_area) * 100 |
