bigwiz/new-plate
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fixed video crash issue

Browse source
This commit is contained in:
Amit Nandi 2025-03-08 20:50:54 +05:30
parent f985128900
commit cd5f375854
2 changed files with 18 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
main.py
View file

@ -31,7 +31,7 @@ CAMERA_SOURCE = 1 #Webcam / Camera Source REMEMBER WINDOWS HAS WEBCAM ACCESS ISS
# Load YOLO models # Load YOLO models
motorcycle_model = YOLO("models/yolov8n.pt") # Pretrained COCO model (motorcycles) motorcycle_model = YOLO("models/yolov8n.pt") # Pretrained COCO model (motorcycles)
helmet_model = YOLO("models/helmetYoloV8_25epochs.pt") # Custom helmet detection model helmet_model = YOLO("models/helmetYoloV8_25epochs.pt") # Custom helmet detection model
plate_model = YOLO("models/license_plate_detector.pt") # Custom license plate detection model plate_model = YOLO("models/license_plate_detector1.pt") # Custom license plate detection model
# Initialize OCR # Initialize OCR
ocr = PaddleOCR(use_angle_cls=True, lang='en') # Angle correction enabled ocr = PaddleOCR(use_angle_cls=True, lang='en') # Angle correction enabled
@ -39,7 +39,7 @@ ocr = PaddleOCR(use_angle_cls=True, lang='en') # Angle correction enabled
# Define Colors # Define Colors
COLOR_MOTORCYCLE = (0, 255, 0) # Green COLOR_MOTORCYCLE = (0, 255, 0) # Green
COLOR_NO_HELMET = (0, 0, 255) # Red COLOR_NO_HELMET = (0, 0, 255) # Red
COLOR_LICENSE_PLATE = (0, 0, 255) # Red COLOR_LICENSE_PLATE = (255, 0, 0) # Blue
#Report Filepath #Report Filepath
EXCEL_FILE = "report.xlsx" EXCEL_FILE = "report.xlsx"
@ -173,7 +173,7 @@ class RealTimeProcessor(QThread):
break break
# Process the frame # Process the frame
processed_frame, _ = self.parent.process_image(frame, is_video=True) processed_frame, license_plate_text = self.parent.process_image(frame, is_video=True)
# Emit signal to update UI # Emit signal to update UI
self.frame_processed.emit(processed_frame) self.frame_processed.emit(processed_frame)
@ -579,6 +579,7 @@ class HelmetDetectionApp(QWidget):
# Format the Excel file # Format the Excel file
format_excel(EXCEL_FILE) format_excel(EXCEL_FILE)
def process_image(self, image_path, is_video=False, output_folder="output/"): def process_image(self, image_path, is_video=False, output_folder="output/"):
# 🖼️ Handle video frames directly # 🖼️ Handle video frames directly
@ -681,12 +682,21 @@ class HelmetDetectionApp(QWidget):
# Convert BGR to RGB (PaddleOCR expects RGB) # Convert BGR to RGB (PaddleOCR expects RGB)
plate_roi_rgb = cv2.cvtColor(plate_roi_denoised, cv2.COLOR_GRAY2RGB) plate_roi_rgb = cv2.cvtColor(plate_roi_denoised, cv2.COLOR_GRAY2RGB)
# Perform OCR with PaddleOCR # Perform OCR with PaddleOCR
result = ocr.ocr(plate_roi_rgb, cls=True) result = ocr.ocr(plate_roi_rgb, cls=True)
print("Raw Result:", result) print("Raw Result:", result)
# Extract text from PaddleOCR result # Perform OCR with PaddleOCR
if result and len(result[0]) > 0: result = ocr.ocr(plate_roi_rgb, cls=True)
print("Raw OCR Result:", result)
# Ensure result is valid
if result is None or not isinstance(result, list) or len(result) == 0:
print("⚠️ OCR result is None or empty!")
license_plate_text = "N/A" # Set default value
elif isinstance(result[0], list) and len(result[0]) > 0:
print("✅ OCR detected text.")
license_plate_text = " ".join([entry[1][0] for entry in result[0]]) # Extract text license_plate_text = " ".join([entry[1][0] for entry in result[0]]) # Extract text
print("🚗 Detected Plate:", license_plate_text) print("🚗 Detected Plate:", license_plate_text)
@ -697,7 +707,9 @@ class HelmetDetectionApp(QWidget):
cv2.rectangle(image, (mx1 + px1, my1 + py1), (mx1 + px2, my1 + py2), COLOR_LICENSE_PLATE, 2) cv2.rectangle(image, (mx1 + px1, my1 + py1), (mx1 + px2, my1 + py2), COLOR_LICENSE_PLATE, 2)
cv2.putText(image, f"Plate: {license_plate_text}", (mx1 + px1, my1 + py1 - 10), cv2.putText(image, f"Plate: {license_plate_text}", (mx1 + px1, my1 + py1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_LICENSE_PLATE, 2) cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_LICENSE_PLATE, 2)
else:
print("⚠️ OCR result does not contain expected text format!")
license_plate_text = "N/A" # Set default value
# 📁 Save processed image for single-image mode # 📁 Save processed image for single-image mode
os.makedirs(output_folder, exist_ok=True) os.makedirs(output_folder, exist_ok=True)

BIN
models/helmetV5_1.pt → models/license_plate_detector1.pt
View file

Binary file not shown.