Abs6187's picture
Update app.py
e66d73b verified
raw
history blame
6.66 kB
import os
import sys
import cv2
import gradio as gr
import numpy as np
import logging
from datetime import datetime
from pathlib import Path
#Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('app.log'),
logging.StreamHandler()
]
)
logger = logging.getLogger(name)
#Add the project root to Python path
project_root = os.path.dirname(os.path.abspath(file))
sys.path.append(project_root)
#Import custom modules and models
from ANPR_IND.scripts.charExtraction import CharExtraction
from ANPR_IND.scripts.bboxAnnotator import BBOXAnnotator
from ultralytics import YOLO
#Initialize ANPR models and classes
wPathPlat = os.path.join(project_root, "ANPR_IND", "licence_plat.pt")
wPathChar = os.path.join(project_root, "ANPR_IND", "licence_character.pt")
classList = np.array([
'A','B','C','D','E','F','G','H','I','J','K','L','M',
'N','O','P','Q','R','S','T','U','V','W','X','Y','Z',
'0','1','2','3','4','5','6','7','8','9'
])
sizePlat = (416, 200)
#Initialize Helmet Detection model path
helmet_model_path = os.path.join(project_root, "Helmet-Detect-model", "best.pt")
#Verify that the required model files exist
required_files = [wPathPlat, wPathChar, helmet_model_path]
for file_path in required_files:
if not os.path.exists(file_path):
logger.error(f"Required model file not found: {file_path}")
raise FileNotFoundError(f"Required model file not found: {file_path}")
#Initialize models
try:
logger.info("Initializing models...")
helmet_model = YOLO(helmet_model_path)
extractor = CharExtraction(
wPlatePath=wPathPlat,
wCharacterPath=wPathChar,
classList=classList,
sizePlate=sizePlat,
conf=0.5
)
annotator = BBOXAnnotator()
logger.info("Models initialized successfully")
except Exception as e:
logger.error(f"Error initializing models: {str(e)}")
raise
def process_image(image, conf=0.45):
start_time = datetime.now()
logger.info(f"Processing image with confidence threshold: {conf}")
text
if image is None:
logger.warning("No image provided")
return None, "No image provided", "No image provided"
try:
# Convert PIL Image to OpenCV BGR format if necessary
if isinstance(image, str):
if not os.path.exists(image):
raise FileNotFoundError(f"Image file not found: {image}")
image = cv2.imread(image)
if image is None:
raise ValueError("Failed to read image from the provided path.")
else:
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
# Run ANPR detection
logger.info("Running ANPR detection")
bbox, plateNum, confidence = extractor.predict(image=image, conf=conf)
anpr_image, plateNum = annotator.draw_bbox(image.copy(), bbox, plateNum)
plate_text = ", ".join(plateNum) if plateNum else "No plate detected"
logger.info(f"ANPR result: {plate_text}")
# Run Helmet detection
logger.info("Running helmet detection")
results = helmet_model(image)
# Ensure accessing the correct results container; the first element usually holds the detection info
helmet_detected = len(results.boxes) > 0
helmet_status = "Helmet Detected" if helmet_detected else "No Helmet Detected"
logger.info(f"Helmet detection result: {helmet_status}")
# Retrieve annotated image from helmet detection
helmet_image = results.plot()
# Combine annotations from both detections
try:
combined_image = cv2.addWeighted(anpr_image, 0.5, helmet_image, 0.5, 0)
except Exception as e:
logger.warning(f"Failed to combine annotations: {str(e)}")
combined_image = helmet_image
# Convert image from BGR to RGB for proper display in Gradio
if isinstance(combined_image, np.ndarray):
combined_image = cv2.cvtColor(combined_image, cv2.COLOR_BGR2RGB)
processing_time = (datetime.now() - start_time).total_seconds()
logger.info(f"Processing completed in {processing_time:.2f} seconds")
return combined_image, plate_text, helmet_status
except Exception as e:
logger.error(f"Error processing image: {str(e)}")
return image, f"Error: {str(e)}", "Error processing image"
#Create an array of example image paths
example_images = [
os.path.join(project_root, "ANPR_IND", "sample_image2.jpg"),
os.path.join(project_root, "ANPR_IND", "sample_image3.jpg"),
os.path.join(project_root, "ANPR_IND", "sample_image5.jpg"),
os.path.join(project_root, "ANPR_IND", "sample_image6.jpg")
]
#Verify example images exist, and remove any that aren't found
for img_path in example_images.copy():
if not os.path.exists(img_path):
logger.warning(f"Example image not found: {img_path}")
example_images.remove(img_path)
def create_interface():
with gr.Blocks(title="Traffic Violation Detection System", theme=gr.themes.Soft()) as demo:
gr.Markdown("# Combined ANPR and Helmet Detection System")
gr.Markdown("Upload an image to detect license plates and check for helmet usage.")
text
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Input Image", type="pil")
conf_slider = gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.45,
label="Confidence Threshold"
)
detect_button = gr.Button("Detect", variant="primary")
with gr.Column():
output_image = gr.Image(label="Annotated Image")
plate_output = gr.Textbox(label="License Plate")
helmet_output = gr.Textbox(label="Helmet Status")
# Configure example images if available
if example_images:
gr.Examples(
examples=[[img, 0.45] for img in example_images],
inputs=[input_image, conf_slider],
outputs=[output_image, plate_output, helmet_output],
fn=process_image,
cache_examples=True
)
# Set up the click event to trigger detection
detect_button.click(
fn=process_image,
inputs=[input_image, conf_slider],
outputs=[output_image, plate_output, helmet_output]
)
return demo
if name == "main":
try:
logger.info("Starting application...")
demo = create_interface()
demo.queue() # Enable request queuing if your tasks are long-running
text
# Launch application on Hugging Face Spaces
# - Use server_name="0.0.0.0" to allow external connections.
# - Do not include share=True (Spaces already provides your public URL).
demo.launch(
server_name="0.0.0.0",
server_port=7860,
debug=True
)
except Exception as e:
logger.error(f"Failed to start application: {str(e)}")
sys.exit(1)