premdeep09/ANPR-System
0
1import cv22from ultralytics import YOLO3import time4 5class ANPRDetector:6 def __init__(self, mode="two_stage", vehicle_model_path="yolov8n.pt", plate_model_path="best_plate.pt"):7 """8 Initializes the ANPR Bounding Box Detector.9 10 Args:11 mode (str): "single_model" (detects both vehicles and plates) or 12 "two_stage" (detects vehicles, then crops and detects plates).13 vehicle_model_path (str): Path to YOLOv8 model for vehicles.14 plate_model_path (str): Path to YOLOv8 model for plates (needed for two-stage, or single-model if it handles both).15 """16 self.mode = mode17 18 if self.mode == "single_model":19 # In single model mode, one YOLO model is trained to detect classes: 0: vehicle, 1: license_plate20 print(f"Loading Single Model from {plate_model_path}...")21 self.model = YOLO(plate_model_path)22 else:23 # Two-stage mode: model 1 detects cars, model 2 detects plates within crops24 print(f"Loading Vehicle Model from {vehicle_model_path}...")25 self.vehicle_model = YOLO(vehicle_model_path)26 27 print(f"Loading Plate Model from {plate_model_path}...")28 # Note: For this demo, assuming you have a trained YOLOv8 for plates. 29 # Fallback to YOLOv8n if file doesn't exist, though it won't detect plates without training.30 try:31 self.plate_model = YOLO(plate_model_path)32 except Exception:33 print(f"Warning: {plate_model_path} not found. Using yolov8n.pt as placeholder.")34 self.plate_model = YOLO("yolov8n.pt")35 36 # Standard COCO classes for vehicles (car, motorcycle, bus, truck)37 self.vehicle_classes = [2, 3, 5, 7]38 39 def process_frame(self, frame):40 """41 Processes a single frame, drawing tight bounding boxes and extracting plate crops.42 """43 processed_frame = frame.copy()44 plate_crops = []45 46 if self.mode == "single_model":47 # Single forward pass for both vehicles and plates48 results = self.model(processed_frame, verbose=False)49 50 for r in results:51 for box in r.boxes:52 cls_id = int(box.cls[0])53 conf = float(box.conf[0])54 x1, y1, x2, y2 = map(int, box.xyxy[0])55 56 if conf < 0.5:57 continue58 59 # Assuming class 0 is Vehicle and class 1 is Plate60 if cls_id == 0:61 # Draw vehicle bounding box62 cv2.rectangle(processed_frame, (x1, y1), (x2, y2), (255, 0, 0), 2)63 cv2.putText(processed_frame, f"Vehicle {conf:.2f}", (x1, max(y1 - 10, 0)), 64 cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)65 elif cls_id == 1:66 # Draw plate bounding box67 cv2.rectangle(processed_frame, (x1, y1), (x2, y2), (0, 255, 255), 2)68 cv2.putText(processed_frame, f"Plate {conf:.2f}", (x1, max(y1 - 10, 0)), 69 cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)70 71 # Crop the plate region72 plate_crop = frame[y1:y2, x1:x2]73 if plate_crop.size > 0:74 plate_crops.append(plate_crop)75 76 elif self.mode == "two_stage":77 # Stage 1: Detect Vehicles78 vehicle_results = self.vehicle_model(processed_frame, classes=self.vehicle_classes, verbose=False)79 80 for r in vehicle_results:81 for box in r.boxes:82 conf = float(box.conf[0])83 if conf < 0.5:84 continue85 86 v_x1, v_y1, v_x2, v_y2 = map(int, box.xyxy[0])87 88 # Draw vehicle bounding box89 cv2.rectangle(processed_frame, (v_x1, v_y1), (v_x2, v_y2), (255, 0, 0), 2)90 cv2.putText(processed_frame, f"Vehicle {conf:.2f}", (v_x1, max(v_y1 - 10, 0)), 91 cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)92 93 vehicle_crop = frame[v_y1:v_y2, v_x1:v_x2]94 if vehicle_crop.size == 0:95 continue96 97 # Stage 2: Detect Plate inside the Vehicle Crop98 plate_results = self.plate_model(vehicle_crop, verbose=False)99 100 for pr in plate_results:101 for p_box in pr.boxes:102 p_conf = float(p_box.conf[0])103 if p_conf < 0.5:104 continue105 106 # Coordinates relative to the crop107 px1, py1, px2, py2 = map(int, p_box.xyxy[0])108 109 # Convert to absolute coordinates110 abs_x1 = v_x1 + px1111 abs_y1 = v_y1 + py1112 abs_x2 = v_x1 + px2113 abs_y2 = v_y1 + py2114 115 # Draw plate bounding box116 cv2.rectangle(processed_frame, (abs_x1, abs_y1), (abs_x2, abs_y2), (0, 255, 255), 2)117 cv2.putText(processed_frame, f"Plate {p_conf:.2f}", (abs_x1, max(abs_y1 - 10, 0)), 118 cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)119 120 # Extract tight crop121 plate_crop = frame[abs_y1:abs_y2, abs_x1:abs_x2]122 if plate_crop.size > 0:123 plate_crops.append(plate_crop)124 125 return processed_frame, plate_crops126 127if __name__ == "__main__":128 # Choose your approach here: "single_model" or "two_stage"129 detector = ANPRDetector(mode="two_stage")130 131 # Optional testing logic132 # cap = cv2.VideoCapture(0)133 # while True:134 # ret, frame = cap.read()135 # if not ret: break136 137 # start_time = time.time()138 # output_frame, crops = detector.process_frame(frame)139 # fps = 1.0 / (time.time() - start_time)140 141 # cv2.putText(output_frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)142 # cv2.imshow("Detection Bounding Boxes", output_frame)143 144 # if cv2.waitKey(1) & 0xFF == ord('q'):145 # break146 # cap.release()147 # cv2.destroyAllWindows()148