utils/annotation.py

from typing import Optional, List
import uuid
from abc import ABC, abstractmethod

class Annotation(ABC):
    box: Optional[List[float]] = None
    index: uuid.UUID = None

    def __init__(self, box: Optional[List[float]] = None):
        self.box = box
        self.index = uuid.uuid4()

class WordAnnotation(Annotation):

    # Class attributes since all classes should have it 
    box:Optional[List[float]] =None,
    score:Optional[float]=None,
    text:str = None, 
    index: uuid.UUID =None

    def __init__(self,
                 box:Optional[List[float]],
                 text:str=None):
        self.box =box if box is not None else None 
        self.text = text if text is not None else None 
        self.index = uuid.uuid4()

class LineAnnotation(Annotation):
    """
    Detection results of all OCR Components 
    `pdf_id` : id or name of pdf so that it can be get from database
    `page`:   pdf page number
    `box`: [xmin, ymin, xmax, ymax]
    
    `index`: same as index for bounding boxes, just the results is wrapped around in this class 
    `score`: prediction score
    `line` : Parent line 
    `text`: text string. 
    """
    # Class attributes since all classes should have it 
    box:Optional[List[float]] =None,
    words:Optional[List[WordAnnotation]] =None,
    index: uuid.UUID =None

    def __init__(self,
                 box:Optional[List[float]],
                 words:List[WordAnnotation]=None):
        self.box = box if box is not None else None 
        self.words = words if words is not None else []
        self.index = uuid.uuid4()