anpr-yolov7 / app.py
lohitkavuru14's picture
Duplicate from itsyoboieltr/anpr-yolov7
9feb5f5
import fileinput
import itertools
import os
import re
from copy import deepcopy
from operator import itemgetter
from pathlib import Path
from typing import Union
import cv2 # type: ignore
import gradio as gr # type: ignore
import numpy as np
import torch
from deep_sort_realtime.deepsort_tracker import DeepSort # type: ignore
from paddleocr import PaddleOCR # type: ignore
if not os.path.isfile("weights.pt"):
weights_url = "https://archive.org/download/anpr_weights/weights.pt"
os.system(f"wget {weights_url}")
if not os.path.isdir("examples"):
examples_url = "https://archive.org/download/anpr_examples_202208/examples.tar.gz"
os.system(f"wget {examples_url}")
os.system("tar -xvf examples.tar.gz")
os.system("rm -rf examples.tar.gz")
def prepend_text(filename: Union[str, Path], text: str):
with fileinput.input(filename, inplace=True) as file:
for line in file:
if file.isfirstline():
print(text)
print(line, end="")
if not os.path.isdir("yolov7"):
yolov7_repo_url = "https://github.com/WongKinYiu/yolov7"
os.system(f"git clone {yolov7_repo_url}")
# Fix import errors
for file in [
"yolov7/models/common.py",
"yolov7/models/experimental.py",
"yolov7/models/yolo.py",
"yolov7/utils/datasets.py",
]:
prepend_text(file, "import sys\nsys.path.insert(0, './yolov7')")
from yolov7.models.experimental import attempt_load # type: ignore
from yolov7.utils.datasets import letterbox # type: ignore
from yolov7.utils.general import check_img_size # type: ignore
from yolov7.utils.general import non_max_suppression # type: ignore
from yolov7.utils.general import scale_coords # type: ignore
from yolov7.utils.plots import plot_one_box # type: ignore
from yolov7.utils.torch_utils import TracedModel, select_device # type: ignore
weights = "weights.pt"
device_id = "cpu"
image_size = 640
trace = True
# Initialize
device = select_device(device_id)
half = device.type != "cpu" # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(image_size, s=stride) # check img_size
if trace:
model = TracedModel(model, device, image_size)
if half:
model.half() # to FP16
if device.type != "cpu":
model(
torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))
) # run once
model.eval()
# Load OCR
paddle = PaddleOCR(lang="en")
def detect_plate(source_image):
# Padded resize
img_size = 640
stride = 32
img = letterbox(source_image, img_size, stride=stride)[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
with torch.no_grad():
# Inference
pred = model(img, augment=True)[0]
# Apply NMS
pred = non_max_suppression(pred, 0.25, 0.45, classes=0, agnostic=True)
plate_detections = []
det_confidences = []
# Process detections
for i, det in enumerate(pred): # detections per image
if len(det):
# Rescale boxes from img_size to source image size
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], source_image.shape
).round()
# Return results
for *xyxy, conf, cls in reversed(det):
coords = [
int(position)
for position in (torch.tensor(xyxy).view(1, 4)).tolist()[0]
]
plate_detections.append(coords)
det_confidences.append(conf.item())
return plate_detections, det_confidences
def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=2.0, threshold=0):
blurred = cv2.GaussianBlur(image, kernel_size, sigma)
sharpened = float(amount + 1) * image - float(amount) * blurred
sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))
sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))
sharpened = sharpened.round().astype(np.uint8)
if threshold > 0:
low_contrast_mask = np.absolute(image - blurred) < threshold
np.copyto(sharpened, image, where=low_contrast_mask)
return sharpened
def crop(image, coord):
cropped_image = image[int(coord[1]) : int(coord[3]), int(coord[0]) : int(coord[2])]
return cropped_image
def ocr_plate(plate_region):
# Image pre-processing for more accurate OCR
rescaled = cv2.resize(
plate_region, None, fx=1.2, fy=1.2, interpolation=cv2.INTER_CUBIC
)
grayscale = cv2.cvtColor(rescaled, cv2.COLOR_BGR2GRAY)
kernel = np.ones((1, 1), np.uint8)
dilated = cv2.dilate(grayscale, kernel, iterations=1)
eroded = cv2.erode(dilated, kernel, iterations=1)
sharpened = unsharp_mask(eroded)
# OCR the preprocessed image
results = paddle.ocr(sharpened, det=False, cls=False)
flattened = list(itertools.chain.from_iterable(results))
plate_text, ocr_confidence = max(flattened, key=itemgetter(1), default=("", 0))
# Filter out anything but uppercase letters, digits, hypens and whitespace.
plate_text = re.sub(r"[^-A-Z0-9 ]", r"", plate_text).strip()
if ocr_confidence == "nan":
ocr_confidence = 0
return plate_text, ocr_confidence
def get_plates_from_image(input):
if input is None:
return None
plate_detections, det_confidences = detect_plate(input)
plate_texts = []
ocr_confidences = []
detected_image = deepcopy(input)
for coords in plate_detections:
plate_region = crop(input, coords)
plate_text, ocr_confidence = ocr_plate(plate_region)
if ocr_confidence == 0: # If OCR confidence is 0, skip this detection
continue
plate_texts.append(plate_text)
ocr_confidences.append(ocr_confidence)
plot_one_box(
coords,
detected_image,
label=plate_text,
color=[0, 150, 255],
line_thickness=2,
)
return detected_image
def pascal_voc_to_coco(x1y1x2y2):
x1, y1, x2, y2 = x1y1x2y2
return [x1, y1, x2 - x1, y2 - y1]
def get_best_ocr(preds, rec_conf, ocr_res, track_id):
for info in preds:
# Check if it is current track id
if info["track_id"] == track_id:
# Check if the ocr confidenence is maximum or not
if info["ocr_conf"] < rec_conf:
info["ocr_conf"] = rec_conf
info["ocr_txt"] = ocr_res
else:
rec_conf = info["ocr_conf"]
ocr_res = info["ocr_txt"]
break
return preds, rec_conf, ocr_res
def get_plates_from_video(source):
if source is None:
return None
# Create a VideoCapture object
video = cv2.VideoCapture(source)
# Default resolutions of the frame are obtained. The default resolutions are system dependent.
# We convert the resolutions from float to integer.
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = video.get(cv2.CAP_PROP_FPS)
# Define the codec and create VideoWriter object.
temp = f"{Path(source).stem}_temp{Path(source).suffix}"
export = cv2.VideoWriter(
temp, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
)
# Intializing tracker
tracker = DeepSort(embedder_gpu=False)
# Initializing some helper variables.
preds = []
total_obj = 0
while True:
ret, frame = video.read()
if ret == True:
# Run the ANPR algorithm
bboxes, scores = detect_plate(frame)
# Convert Pascal VOC detections to COCO
bboxes = list(map(lambda bbox: pascal_voc_to_coco(bbox), bboxes))
if len(bboxes) > 0:
# Storing all the required info in a list.
detections = [
(bbox, score, "number_plate") for bbox, score in zip(bboxes, scores)
]
# Applying tracker.
# The tracker code flow: kalman filter -> target association(using hungarian algorithm) and appearance descriptor.
tracks = tracker.update_tracks(detections, frame=frame)
# Checking if tracks exist.
for track in tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
# Changing track bbox to top left, bottom right coordinates
bbox = [int(position) for position in list(track.to_tlbr())]
for i in range(len(bbox)):
if bbox[i] < 0:
bbox[i] = 0
# Cropping the license plate and applying the OCR.
plate_region = crop(frame, bbox)
plate_text, ocr_confidence = ocr_plate(plate_region)
# Storing the ocr output for corresponding track id.
output_frame = {
"track_id": track.track_id,
"ocr_txt": plate_text,
"ocr_conf": ocr_confidence,
}
# Appending track_id to list only if it does not exist in the list
# else looking for the current track in the list and updating the highest confidence of it.
if track.track_id not in list(
set(pred["track_id"] for pred in preds)
):
total_obj += 1
preds.append(output_frame)
else:
preds, ocr_confidence, plate_text = get_best_ocr(
preds, ocr_confidence, plate_text, track.track_id
)
# Plotting the prediction.
plot_one_box(
bbox,
frame,
label=f"{str(track.track_id)}. {plate_text}",
color=[255, 150, 0],
line_thickness=3,
)
# Write the frame into the output file
export.write(frame)
else:
break
# When everything done, release the video capture and video write objects
video.release()
export.release()
# Compressing the output video for smaller size and web compatibility.
output = f"{Path(source).stem}_detected{Path(source).suffix}"
os.system(
f"ffmpeg -y -i {temp} -c:v libx264 -b:v 5000k -minrate 1000k -maxrate 8000k -pass 1 -c:a aac -f mp4 /dev/null && ffmpeg -i {temp} -c:v libx264 -b:v 5000k -minrate 1000k -maxrate 8000k -pass 2 -c:a aac -movflags faststart {output}"
)
os.system(f"rm -rf {temp} ffmpeg2pass-0.log ffmpeg2pass-0.log.mbtree")
return output
with gr.Blocks() as demo:
gr.Markdown('### <h3 align="center">Automatic Number Plate Recognition</h3>')
gr.Markdown(
"This AI was trained to detect and recognize number plates on vehicles."
)
with gr.Tabs():
with gr.TabItem("Image"):
with gr.Row():
image_input = gr.Image()
image_output = gr.Image()
image_input.change(
get_plates_from_image, inputs=image_input, outputs=image_output
)
gr.Examples(
[
["examples/test_image_1.jpg"],
["examples/test_image_2.jpg"],
["examples/test_image_3.png"],
["examples/test_image_4.jpeg"],
],
[image_input],
image_output,
get_plates_from_image,
cache_examples=True,
)
with gr.TabItem("Video"):
with gr.Row():
video_input = gr.Video(format="mp4")
video_output = gr.Video(format="mp4")
video_input.change(
get_plates_from_video, inputs=video_input, outputs=video_output
)
gr.Examples(
[["examples/test_video_1.mp4"]],
[video_input],
video_output,
get_plates_from_video,
cache_examples=True,
)
gr.Markdown("[@itsyoboieltr](https://github.com/itsyoboieltr)")
demo.launch()