gatekeeper / app.py
vikramjeetthakur's picture
Create app.py
3ae443a verified
raw
history blame
3.03 kB
from transformers import DetrImageProcessor, DetrForObjectDetection, TrOCRProcessor, VisionEncoderDecoderModel
import cv2
from PIL import Image, ImageDraw
import torch
import streamlit as st
# Load Hugging Face Models
detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
# Detect license plates
def detect_license_plate(frame):
pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
inputs = detr_processor(images=pil_image, return_tensors="pt")
outputs = detr_model(**inputs)
target_sizes = torch.tensor([pil_image.size[::-1]])
results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)
return results[0]["boxes"], pil_image
# Recognize text
def recognize_text_from_plate(cropped_plate):
inputs = trocr_processor(images=cropped_plate, return_tensors="pt")
outputs = trocr_model.generate(**inputs)
return trocr_processor.batch_decode(outputs, skip_special_tokens=True)[0]
# Streamlit configuration
st.title("Real-Time Car Number Plate Recognition")
st.text("This application uses Hugging Face Transformers to detect and recognize car plates.")
# Authorized car database
authorized_cars = {"KA01AB1234", "MH12XY5678", "DL8CAF9090"}
# Verification function
def verify_plate(plate_text):
if plate_text in authorized_cars:
return f"βœ… Access Granted: {plate_text}"
else:
return f"❌ Access Denied: {plate_text}"
# Live video feed and processing
def live_feed():
cap = cv2.VideoCapture(0) # Open the webcam
stframe = st.empty() # Streamlit frame for displaying video
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Detect license plates
boxes, pil_image = detect_license_plate(frame)
draw = ImageDraw.Draw(pil_image)
recognized_plates = []
for box in boxes:
# Crop the detected plate
cropped_plate = pil_image.crop((box[0], box[1], box[2], box[3]))
# Recognize text
plate_text = recognize_text_from_plate(cropped_plate)
recognized_plates.append(plate_text)
# Draw bounding box and text
draw.rectangle(box.tolist(), outline="red", width=2)
draw.text((box[0], box[1]), plate_text, fill="red")
# Convert PIL image back to OpenCV format
processed_frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
# Stream the video to Streamlit
stframe.image(processed_frame, channels="BGR")
# Show results
for plate_text in recognized_plates:
st.write(verify_plate(plate_text))
cap.release()
cv2.destroyAllWindows()
if st.button("Start Camera"):
live_feed()