Spaces:
Runtime error
Runtime error
| # 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 | |