anpr.py

import numpy as np
import cv2
import matplotlib.pyplot as plt
#import pytesseract as pt
import os
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import requests
import torch
import plotly.express as px

# LOAD YOLO MODEL
INPUT_WIDTH =  640
INPUT_HEIGHT = 640
#onnx_file_path = os.path.abspath('./static/models/best.onnx')
onnx_file_path = os.path.abspath(r'./static/model/best.onnx')
print(f"Attempting to load ONNX file from: {onnx_file_path}")
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-small-printed')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-small-printed') 



if not os.path.exists(onnx_file_path):
    print(f"Error: ONNX file not found at {onnx_file_path}")
else:
    try:
        net = cv2.dnn.readNetFromONNX(onnx_file_path)
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
    except cv2.error as e:
        print(f"Error loading ONNX file: {onnx_file_path}")
        print(f"OpenCV error: {e}")
        raise  # Re-raise the exception to halt the program





def get_detections(img,net):
    # CONVERT IMAGE TO YOLO FORMAT
    image = img.copy()
    row, col, d = image.shape

    max_rc = max(row,col)
    input_image = np.zeros((max_rc,max_rc,3),dtype=np.uint8)
    input_image[0:row,0:col] = image

    # GET PREDICTION FROM YOLO MODEL
    blob = cv2.dnn.blobFromImage(input_image,1/255,(INPUT_WIDTH,INPUT_HEIGHT),swapRB=True,crop=False)
    net.setInput(blob)
    preds = net.forward()
    detections = preds[0]
    
    return input_image, detections

def non_maximum_supression(input_image,detections):
    # FILTER DETECTIONS BASED ON CONFIDENCE AND PROBABILIY SCORE
    # center x, center y, w , h, conf, proba
    boxes = []
    confidences = []

    image_w, image_h = input_image.shape[:2]
    x_factor = image_w/INPUT_WIDTH
    y_factor = image_h/INPUT_HEIGHT

    for i in range(len(detections)):
        row = detections[i]
        confidence = row[4] # confidence of detecting license plate
        if confidence > 0.4:
            class_score = row[5] # probability score of license plate
            if class_score > 0.25:
                cx, cy , w, h = row[0:4]

                left = int((cx - 0.5*w)*x_factor)
                top = int((cy-0.5*h)*y_factor)
                width = int(w*x_factor)
                height = int(h*y_factor)
                box = np.array([left,top,width,height])

                confidences.append(confidence)
                boxes.append(box)

    # clean
    boxes_np = np.array(boxes).tolist()
    confidences_np = np.array(confidences).tolist()
    # NMS
    index = np.array(cv2.dnn.NMSBoxes(boxes_np,confidences_np,0.25,0.45)).flatten()
    
    return boxes_np, confidences_np, index

def extract_text_py(image,bbox):
    x,y,w,h = bbox
    
    roi = image[y:y+h, x:x+w]
    if 0 in roi.shape:
        return ''
    else:
        roi_bgr = cv2.cvtColor(roi,cv2.COLOR_RGB2BGR)
        gray = cv2.cvtColor(roi_bgr,cv2.COLOR_BGR2GRAY)
        magic_color = apply_brightness_contrast(gray,brightness=40,contrast=70)
        #text = pt.image_to_string(magic_color)
        text = pt.image_to_string(magic_color,lang='eng',config='--psm 6')
        text = text.strip()
        
        return text
    
# extrating text
def extract_text(image,bbox):
    x,y,w,h = bbox
    roi = image[y:y+h, x:x+w]
    #print("roi:",roi)

    # Use OpenCV to read the image
    img = roi.copy()
    print(img.shape)
    # Convert BGR to RGB
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # Create the Plotly Express figure
    #fig = px.imshow(img_rgb)

    # Update layout and show the figure
    #fig.update_layout(width=100, height=40, margin=dict(l=10, r=10, b=10, t=10))
    #fig.update_xaxes(showticklabels=False).update_yaxes(showticklabels=False)
    #fig.show()

    image = img_rgb
    if 0 in roi.shape:
        return 'no number'

    else:
        pixel_values = processor(images=image, return_tensors="pt").pixel_values
        generated_ids = model.generate(pixel_values)
        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        text = filter_string(text)

        return text
    
def filter_string(input_string):
    filtered_chars = [char for char in input_string if char.isalnum() and (char.isupper() or char.isdigit())]
    filtered_string = ''.join(filtered_chars)
    return filtered_string
  

def drawings(image,boxes_np,confidences_np,index):
    # drawings
    text_list = []
    for ind in index:
        x,y,w,h =  boxes_np[ind]
        bb_conf = confidences_np[ind]
        conf_text = 'plate: {:.0f}%'.format(bb_conf*100)
        license_text = extract_text(image,boxes_np[ind])


        cv2.rectangle(image,(x,y),(x+w,y+h),(255,0,255),2)
        cv2.rectangle(image,(x,y-30),(x+w,y),(255,0,255),-1)
        cv2.rectangle(image,(x,y+h),(x+w,y+h+30),(0,0,0),-1)


        cv2.putText(image,conf_text,(x,y-10),cv2.FONT_HERSHEY_SIMPLEX,0.7,(255,255,255),1)
        cv2.putText(image,license_text,(x,y+h+27),cv2.FONT_HERSHEY_SIMPLEX,0.7,(0,255,0),1)
        
        text_list.append(license_text)

    return image,  text_list


# predictions
def yolo_predictions(img,net):
    ## step-1: detections
    input_image, detections = get_detections(img,net)
    ## step-2: NMS
    boxes_np, confidences_np, index = non_maximum_supression(input_image, detections)
    ## step-3: Drawings
    result_img, text = drawings(img,boxes_np,confidences_np,index)
    return result_img, text


def object_detection(path,filename):
    # read image
    image = cv2.imread(path) # PIL object
    image = np.array(image,dtype=np.uint8) # 8 bit array (0,255)
    result_img, text_list = yolo_predictions(image,net)
    cv2.imwrite('./static/predict/{}'.format(filename),result_img)
    return text_list



# def OCR(path,filename):
#     img = np.array(load_img(path))
#     cods = object_detection(path,filename)
#     xmin ,xmax,ymin,ymax = cods[0]
#     roi = img[ymin:ymax,xmin:xmax]
#     roi_bgr = cv2.cvtColor(roi,cv2.COLOR_RGB2BGR)
#     gray = cv2.cvtColor(roi_bgr,cv2.COLOR_BGR2GRAY)
#     magic_color = apply_brightness_contrast(gray,brightness=40,contrast=70)
#     cv2.imwrite('./static/roi/{}'.format(filename),roi_bgr)
    
    
    
#     print(text)
#     save_text(filename,text)
#     return text


def apply_brightness_contrast(input_img, brightness = 0, contrast = 0):
    
        if brightness != 0:
            if brightness > 0:
                shadow = brightness
                highlight = 255
            else:
                shadow = 0
                highlight = 255 + brightness
            alpha_b = (highlight - shadow)/255
            gamma_b = shadow
            
            buf = cv2.addWeighted(input_img, alpha_b, input_img, 0, gamma_b)
        else:
            buf = input_img.copy()
        
        if contrast != 0:
            f = 131*(contrast + 127)/(127*(131-contrast))
            alpha_c = f
            gamma_c = 127*(1-f)
            
            buf = cv2.addWeighted(buf, alpha_c, buf, 0, gamma_c)

        return buf
    #