app.py

import PIL.Image
import gradio as gr
import torch
import numpy as np

dw=0.3
dh=0.25
def is_nw(box):
    """
    A box happen to be a 4-pixel list in order
    1 -- 2
    4 -- 3
    """
    return box[2][0]<=dw and box[2][1]<= dh

def is_ne(box):
    return box[3][0]>=1-dw and box[3][1]<= dh

def is_se(box):
    return box[0][0]>=1-dw and box[0][1]>= 1-dh

def is_sw(box):
    return box[1][0]<=dw and box[1][1]>= 1-dh

def is_corner(box)->bool:
    """ @:returns true if the box is located in any corner """
    return is_nw(box) or is_ne(box) or is_se(box) or is_sw(box)

dhhf=0.2 # dh for header and footer
def is_footer(box)->bool:
    """ true if for the 2 first points, y>0.8 """
    return box[0][1]>=1-dhhf and box[1][1]>=1-dhhf

def is_header(box)->bool:
    """ true if for the 2 last points, y<0.2 """
    return box[2][1]<=dhhf and box[3][1]<=dhhf

def is_signature(prediction_result) -> bool:
    """ true if any of the boxes is at any corner """
    for box in prediction_result['boxes_as_ratios']:
        if is_corner(box) or is_header(box) or is_footer(box):
            return True
    return False

def detect_with_craft_text_detector(image: PIL.Image.Image):
    from craft_text_detector import Craft
    craft = Craft(output_dir='output', crop_type="box", cuda=torch.cuda.is_available(), export_extra=True)
    result = craft.detect_text( np.asarray(image))
    annotated = PIL.Image.open('output/image_text_detection.png')  # image with boxes displayed
    return annotated, result['boxes'], is_signature(result)

def detect_with_craft_hw_ocr(image: PIL.Image.Image):
    image = np.asarray(image)
    from craft_hw_ocr import OCR
    ocr = OCR.load_models()
    image, results = OCR.detection(image, ocr[2])
    bboxes, _ = OCR.recoginition(image, results, ocr[0], ocr[1])
    annotated = OCR.visualize(image, results)
    return annotated, bboxes, False

def process(image:PIL.Image.Image, lib:str):
    if image is None:
        return None,0,False
    annotated, boxes, signed = detect_with_craft_text_detector(image) if lib=='craft_text_detector' else detect_with_craft_hw_ocr( image)
    return annotated, len(boxes), signed

gr.Interface(
    fn = process,
    inputs = [ gr.Image(type="pil", label="Input"), gr.inputs.Radio(label='Model', choices=["craft_text_detector", "craft_hw_ocr"], default='craft_text_detector') ],
    outputs = [ gr.Image(type="pil", label="Output"), gr.Label(label="nb of text detections"), gr.Label(label="Has signature") ],
    title="Detect signature in image",
    description="Is the photo or image watermarked by a signature?",
    examples=[['data/photologo-1-1.jpg'], ['data/times-square.jpg']],
    allow_flagging="never"
).launch(debug=True, enable_queue=True)