Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -3,62 +3,97 @@ from streamlit_webrtc import webrtc_streamer, VideoProcessorBase
|
|
3 |
import av
|
4 |
from transformers import DetrImageProcessor, DetrForObjectDetection, TrOCRProcessor, VisionEncoderDecoderModel
|
5 |
from PIL import Image, ImageDraw
|
6 |
-
import torch
|
7 |
import numpy as np
|
|
|
8 |
|
9 |
-
# Load Models
|
|
|
10 |
detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
|
11 |
detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
|
|
|
|
|
12 |
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
|
13 |
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
|
14 |
|
15 |
-
# Authorized car database
|
16 |
-
authorized_cars = {"KA01AB1234", "MH12XY5678", "DL8CAF9090"}
|
17 |
|
18 |
-
|
|
|
19 |
def detect_license_plate(frame):
|
|
|
|
|
|
|
20 |
pil_image = Image.fromarray(frame)
|
21 |
inputs = detr_processor(images=pil_image, return_tensors="pt")
|
22 |
outputs = detr_model(**inputs)
|
|
|
|
|
23 |
target_sizes = torch.tensor([pil_image.size[::-1]])
|
24 |
results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)
|
25 |
return results[0]["boxes"], pil_image
|
26 |
|
27 |
-
|
28 |
def recognize_text_from_plate(cropped_plate):
|
|
|
|
|
|
|
29 |
inputs = trocr_processor(images=cropped_plate, return_tensors="pt")
|
30 |
outputs = trocr_model.generate(**inputs)
|
31 |
return trocr_processor.batch_decode(outputs, skip_special_tokens=True)[0]
|
32 |
|
33 |
-
|
34 |
def verify_plate(plate_text):
|
|
|
|
|
|
|
35 |
if plate_text in authorized_cars:
|
36 |
return f"✅ Access Granted: {plate_text}"
|
37 |
else:
|
38 |
return f"❌ Access Denied: {plate_text}"
|
39 |
|
40 |
-
|
|
|
41 |
class LicensePlateProcessor(VideoProcessorBase):
|
|
|
|
|
|
|
42 |
def recv(self, frame: av.VideoFrame):
|
43 |
-
frame = frame.to_ndarray(format="bgr24")
|
44 |
boxes, pil_image = detect_license_plate(frame)
|
45 |
draw = ImageDraw.Draw(pil_image)
|
46 |
|
47 |
recognized_plates = []
|
48 |
for box in boxes:
|
|
|
49 |
cropped_plate = pil_image.crop((box[0], box[1], box[2], box[3]))
|
50 |
plate_text = recognize_text_from_plate(cropped_plate)
|
51 |
recognized_plates.append(plate_text)
|
|
|
|
|
52 |
draw.rectangle(box.tolist(), outline="red", width=3)
|
53 |
draw.text((box[0], box[1]), plate_text, fill="red")
|
54 |
|
55 |
-
#
|
56 |
processed_frame = np.array(pil_image)
|
|
|
|
|
57 |
for plate_text in recognized_plates:
|
58 |
st.write(verify_plate(plate_text))
|
|
|
59 |
return av.VideoFrame.from_ndarray(processed_frame, format="bgr24")
|
60 |
|
61 |
-
|
|
|
62 |
st.title("Real-Time Car Number Plate Recognition")
|
63 |
-
st.write("
|
64 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3 |
import av
|
4 |
from transformers import DetrImageProcessor, DetrForObjectDetection, TrOCRProcessor, VisionEncoderDecoderModel
|
5 |
from PIL import Image, ImageDraw
|
|
|
6 |
import numpy as np
|
7 |
+
import torch
|
8 |
|
9 |
+
# Step 1: Load Models
|
10 |
+
# DETR for object detection
|
11 |
detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
|
12 |
detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
|
13 |
+
|
14 |
+
# TrOCR for text recognition
|
15 |
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
|
16 |
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
|
17 |
|
18 |
+
# Authorized car database for verification
|
19 |
+
authorized_cars = {"KA01AB1234", "MH12XY5678", "DL8CAF9090"} # Example data
|
20 |
|
21 |
+
|
22 |
+
# Step 2: Define Helper Functions
|
23 |
def detect_license_plate(frame):
|
24 |
+
"""
|
25 |
+
Detect license plates in the frame using DETR.
|
26 |
+
"""
|
27 |
pil_image = Image.fromarray(frame)
|
28 |
inputs = detr_processor(images=pil_image, return_tensors="pt")
|
29 |
outputs = detr_model(**inputs)
|
30 |
+
|
31 |
+
# Get bounding boxes
|
32 |
target_sizes = torch.tensor([pil_image.size[::-1]])
|
33 |
results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)
|
34 |
return results[0]["boxes"], pil_image
|
35 |
|
36 |
+
|
37 |
def recognize_text_from_plate(cropped_plate):
|
38 |
+
"""
|
39 |
+
Recognize text from the cropped license plate image using TrOCR.
|
40 |
+
"""
|
41 |
inputs = trocr_processor(images=cropped_plate, return_tensors="pt")
|
42 |
outputs = trocr_model.generate(**inputs)
|
43 |
return trocr_processor.batch_decode(outputs, skip_special_tokens=True)[0]
|
44 |
|
45 |
+
|
46 |
def verify_plate(plate_text):
|
47 |
+
"""
|
48 |
+
Check if the recognized plate text exists in the authorized cars database.
|
49 |
+
"""
|
50 |
if plate_text in authorized_cars:
|
51 |
return f"✅ Access Granted: {plate_text}"
|
52 |
else:
|
53 |
return f"❌ Access Denied: {plate_text}"
|
54 |
|
55 |
+
|
56 |
+
# Step 3: Custom Video Processor for WebRTC
|
57 |
class LicensePlateProcessor(VideoProcessorBase):
|
58 |
+
"""
|
59 |
+
Custom video processor to handle video frames in real-time.
|
60 |
+
"""
|
61 |
def recv(self, frame: av.VideoFrame):
|
62 |
+
frame = frame.to_ndarray(format="bgr24") # Convert frame to NumPy array
|
63 |
boxes, pil_image = detect_license_plate(frame)
|
64 |
draw = ImageDraw.Draw(pil_image)
|
65 |
|
66 |
recognized_plates = []
|
67 |
for box in boxes:
|
68 |
+
# Crop detected license plate
|
69 |
cropped_plate = pil_image.crop((box[0], box[1], box[2], box[3]))
|
70 |
plate_text = recognize_text_from_plate(cropped_plate)
|
71 |
recognized_plates.append(plate_text)
|
72 |
+
|
73 |
+
# Draw bounding box and label on the image
|
74 |
draw.rectangle(box.tolist(), outline="red", width=3)
|
75 |
draw.text((box[0], box[1]), plate_text, fill="red")
|
76 |
|
77 |
+
# Convert back to OpenCV format
|
78 |
processed_frame = np.array(pil_image)
|
79 |
+
|
80 |
+
# Log results in Streamlit UI
|
81 |
for plate_text in recognized_plates:
|
82 |
st.write(verify_plate(plate_text))
|
83 |
+
|
84 |
return av.VideoFrame.from_ndarray(processed_frame, format="bgr24")
|
85 |
|
86 |
+
|
87 |
+
# Step 4: Streamlit Interface
|
88 |
st.title("Real-Time Car Number Plate Recognition")
|
89 |
+
st.write("This app uses Hugging Face Transformers and WebRTC for real-time processing.")
|
90 |
+
|
91 |
+
# Start WebRTC Streamer
|
92 |
+
webrtc_streamer(
|
93 |
+
key="plate-recognition",
|
94 |
+
video_processor_factory=LicensePlateProcessor,
|
95 |
+
rtc_configuration={
|
96 |
+
# Required to ensure WebRTC works across networks
|
97 |
+
"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
|
98 |
+
}
|
99 |
+
)
|