detector/utils.py

import torch
import numpy as np
import cv2
import os
from datetime import datetime
from json import dumps
import requests

BASE_DIR = os.path.abspath(os.getcwd())

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Loading models...', device)

model_plates = torch.hub.load('ultralytics/yolov5', 'custom',
                              path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'), device=device)

model_chars = torch.hub.load('ultralytics/yolov5', 'custom',
                             path=os.path.join(BASE_DIR, 'detector', 'static', 'chars.pt'), device=device)


def pad_img_to_fit_bbox(img, x1, x2, y1, y2):
    img = np.pad(img, ((np.abs(np.minimum(0, y1)), np.maximum(y2 - img.shape[0], 0)),
                       (np.abs(np.minimum(0, x1)), np.maximum(x2 - img.shape[1], 0)), (0, 0)), mode="constant")
    y1 += np.abs(np.minimum(0, y1))
    y2 += np.abs(np.minimum(0, y1))
    x1 += np.abs(np.minimum(0, x1))
    x2 += np.abs(np.minimum(0, x1))
    return img, x1, x2, y1, y2


def imcrop(img, bbox):
    x1, y1, x2, y2 = bbox
    if x1 < 0 or y1 < 0 or x2 > img.shape[1] or y2 > img.shape[0]:
        img, x1, x2, y1, y2 = pad_img_to_fit_bbox(img, x1, x2, y1, y2)
    return img[y1:y2, x1:x2, :]


def detect_plates(img):
    return model_plates(img)


def detect_chars(img):
    img = cv2.resize(img, (640, 320))
    detect = model_chars(img)
    records = detect.pandas().xyxy[0].to_dict(orient='records')
    text = ''
    if records:
        records = sorted(records, key=lambda d: d['xmin'])
        text = ''.join([i.get('name') for i in records])
    return text


def draw_text(img, text,
              pos=(0, 0),
              font_scale=1,
              font_thickness=2,
              text_color=(0, 255, 0),
              text_color_bg=(0, 0, 0)
              ):
    x, y = pos
    text_size, _ = cv2.getTextSize(text, 0, font_scale, font_thickness)
    text_w, text_h = text_size
    cv2.rectangle(img, pos, (x + text_w, y - text_h), text_color_bg, -1)
    cv2.putText(img, text, (x, y), 0, font_scale, text_color, font_thickness)


def send_request(frame, text, bbox):
    cv2.rectangle(
        frame,
        (bbox[0], bbox[1]),
        (bbox[2], bbox[3]),
        (0, 255, 0),
        2,
    )
    draw_text(frame, text, (bbox[0], bbox[1]))
    url = "https://api.prevantec.com/toll-plates"
    data = {
        "number": text,
        "camera": "camera_1",
        "spot_on": str(datetime.now()),
    }
    if not os.path.exists(os.path.join(BASE_DIR, 'plates')):
        os.makedirs(os.path.join(BASE_DIR, 'plates'))
    filename = os.path.join(BASE_DIR, 'plates', f'{text}.jpg')
    cv2.imwrite(filename, frame)
    payload = {'data': dumps(data)}
    files = [
        ('files', (f'{text}.jpg', open(filename, 'rb'), 'image/jpg'))
    ]
    headers = {}
    requests.request("POST", url, headers=headers, data=payload, files=files)