File size: 2,912 Bytes
3c84e45 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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)
|