Prabhat51/EdgeAI_Project

1# inference_utils.py2import os, cv2, re3import torch4import pandas as pd5from ultralytics import YOLO6from datetime import datetime7from paddleocr import PaddleOCR8from difflib import get_close_matches9 10from huggingface_hub import hf_hub_download11from torch.serialization import safe_globals12from ultralytics.nn.tasks import DetectionModel13from ultralytics import YOLO14# Download to local cache15 16 17# Load models from Hugging Face18def load_models():19    # vehicle_detector = YOLO("https://huggingface.co/Prabhat51/number-plate-models/blob/main/veh_detect.pt")20    # vehicle_classifier = YOLO("https://huggingface.co/Prabhat51/number-plate-models/blob/main/veh_class.pt")21    # plate_detector = YOLO("https://huggingface.co/Prabhat51/number-plate-models/blob/main/plate_detect.pt")22    veh_detect_path = hf_hub_download(repo_id="Prabhat51/number-plate-models", filename="veh_detect.pt")23    with safe_globals([DetectionModel]):24        vehicle_detector = YOLO(veh_detect_path)25    # vehicle_detector = YOLO(veh_detect_path)26    vehicle_classifier_path = hf_hub_download(repo_id="Prabhat51/number-plate-models", filename="veh_class.pt")27    vehicle_classifier = YOLO(vehicle_classifier_path)28    plate_detector_path = hf_hub_download(repo_id="Prabhat51/number-plate-models", filename="plate_detect.pt")29    with safe_globals([DetectionModel]):30        vehicle_detector = YOLO(plate_detector_path)31    # plate_detector = YOLO(plate_detector_path)32    ocr_reader = PaddleOCR(use_angle_cls=True, lang='en')33    return vehicle_detector, vehicle_classifier, plate_detector, ocr_reader34 35# Validate Indian number plate36valid_rto_codes = { ... }  # use your RTO set here37 38def correct_plate_text(text):39    text = re.sub(r'[^A-Z0-9]', '', text.upper())40    text = text.replace('O', '0').replace('I', '1')41    match = re.match(r'^([A-Z]{2})([0-9]{2})([A-Z]{1,2})([0-9]{3,4})$', text)42    if match and match.group(1) in valid_rto_codes:43        return text44    return None45 46# Inference on single frame47def process_frame(frame, vehicle_detector, vehicle_classifier, plate_detector, ocr_reader):48    results = []49    detections = vehicle_detector(frame)[0].boxes50    for box in detections:51        x1, y1, x2, y2 = map(int, box.xyxy[0])52        vehicle_crop = frame[y1:y2, x1:x2]53 54        cls_result = vehicle_classifier(vehicle_crop)55        if not cls_result[0].probs:56            continue57        vehicle_type = cls_result[0].names[cls_result[0].probs.top1]58 59        plate_boxes = plate_detector(vehicle_crop)[0].boxes60        for pb in plate_boxes:61            px1, py1, px2, py2 = map(int, pb.xyxy[0])62            plate_crop = vehicle_crop[py1:py2, px1:px2]63 64            ocr_result = ocr_reader.ocr(plate_crop, cls=True)65            if not ocr_result or not ocr_result[0]:66                continue67 68            raw_text = ocr_result[0][0][1][0]69            plate_text = correct_plate_text(raw_text)70            if not plate_text:71                continue72 73            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")74            results.append((timestamp, vehicle_type, plate_text))75    return results76