Create app.py

app.py

@@ -1,368 +0,0 @@
-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()