Spaces:
Running
Running
import streamlit as st | |
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase | |
import av | |
from transformers import DetrImageProcessor, DetrForObjectDetection, TrOCRProcessor, VisionEncoderDecoderModel | |
from PIL import Image, ImageDraw | |
import torch | |
import numpy as np | |
# Load 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") | |
# Authorized car database | |
authorized_cars = {"KA01AB1234", "MH12XY5678", "DL8CAF9090"} | |
# Detect License Plates | |
def detect_license_plate(frame): | |
pil_image = Image.fromarray(frame) | |
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 from Plates | |
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] | |
# Verify Plate | |
def verify_plate(plate_text): | |
if plate_text in authorized_cars: | |
return f"β Access Granted: {plate_text}" | |
else: | |
return f"β Access Denied: {plate_text}" | |
# Custom Video Processor | |
class LicensePlateProcessor(VideoProcessorBase): | |
def recv(self, frame: av.VideoFrame): | |
frame = frame.to_ndarray(format="bgr24") | |
boxes, pil_image = detect_license_plate(frame) | |
draw = ImageDraw.Draw(pil_image) | |
recognized_plates = [] | |
for box in boxes: | |
cropped_plate = pil_image.crop((box[0], box[1], box[2], box[3])) | |
plate_text = recognize_text_from_plate(cropped_plate) | |
recognized_plates.append(plate_text) | |
draw.rectangle(box.tolist(), outline="red", width=3) | |
draw.text((box[0], box[1]), plate_text, fill="red") | |
# Return processed frame | |
processed_frame = np.array(pil_image) | |
for plate_text in recognized_plates: | |
st.write(verify_plate(plate_text)) | |
return av.VideoFrame.from_ndarray(processed_frame, format="bgr24") | |
# Streamlit UI | |
st.title("Real-Time Car Number Plate Recognition") | |
st.write("Streamlit with WebRTC for camera streaming.") | |
webrtc_streamer(key="plate-recognition", video_processor_factory=LicensePlateProcessor) | |