Spaces:
Runtime error
Runtime error
File size: 3,098 Bytes
9d5b1af 72c5da0 4acf106 72c5da0 9d5b1af 72c5da0 4acf106 72c5da0 4acf106 9d5b1af |
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 |
# inference_utils.py
import os, cv2, re
import torch
import pandas as pd
from ultralytics import YOLO
from datetime import datetime
from paddleocr import PaddleOCR
from difflib import get_close_matches
from huggingface_hub import hf_hub_download
from torch.serialization import safe_globals
from ultralytics.nn.tasks import DetectionModel
from ultralytics import YOLO
# Download to local cache
# Load models from Hugging Face
def load_models():
# vehicle_detector = YOLO("https://huggingface.co/Prabhat51/number-plate-models/blob/main/veh_detect.pt")
# vehicle_classifier = YOLO("https://huggingface.co/Prabhat51/number-plate-models/blob/main/veh_class.pt")
# plate_detector = YOLO("https://huggingface.co/Prabhat51/number-plate-models/blob/main/plate_detect.pt")
veh_detect_path = hf_hub_download(repo_id="Prabhat51/number-plate-models", filename="veh_detect.pt")
with safe_globals([DetectionModel]):
vehicle_detector = YOLO(veh_detect_path)
# vehicle_detector = YOLO(veh_detect_path)
vehicle_classifier_path = hf_hub_download(repo_id="Prabhat51/number-plate-models", filename="veh_class.pt")
vehicle_classifier = YOLO(vehicle_classifier_path)
plate_detector_path = hf_hub_download(repo_id="Prabhat51/number-plate-models", filename="plate_detect.pt")
with safe_globals([DetectionModel]):
vehicle_detector = YOLO(plate_detector_path)
# plate_detector = YOLO(plate_detector_path)
ocr_reader = PaddleOCR(use_angle_cls=True, lang='en')
return vehicle_detector, vehicle_classifier, plate_detector, ocr_reader
# Validate Indian number plate
valid_rto_codes = { ... } # use your RTO set here
def correct_plate_text(text):
text = re.sub(r'[^A-Z0-9]', '', text.upper())
text = text.replace('O', '0').replace('I', '1')
match = re.match(r'^([A-Z]{2})([0-9]{2})([A-Z]{1,2})([0-9]{3,4})$', text)
if match and match.group(1) in valid_rto_codes:
return text
return None
# Inference on single frame
def process_frame(frame, vehicle_detector, vehicle_classifier, plate_detector, ocr_reader):
results = []
detections = vehicle_detector(frame)[0].boxes
for box in detections:
x1, y1, x2, y2 = map(int, box.xyxy[0])
vehicle_crop = frame[y1:y2, x1:x2]
cls_result = vehicle_classifier(vehicle_crop)
if not cls_result[0].probs:
continue
vehicle_type = cls_result[0].names[cls_result[0].probs.top1]
plate_boxes = plate_detector(vehicle_crop)[0].boxes
for pb in plate_boxes:
px1, py1, px2, py2 = map(int, pb.xyxy[0])
plate_crop = vehicle_crop[py1:py2, px1:px2]
ocr_result = ocr_reader.ocr(plate_crop, cls=True)
if not ocr_result or not ocr_result[0]:
continue
raw_text = ocr_result[0][0][1][0]
plate_text = correct_plate_text(raw_text)
if not plate_text:
continue
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
results.append((timestamp, vehicle_type, plate_text))
return results
|