vikramjeetthakur commited on
Commit
04b96ce
·
verified ·
1 Parent(s): 093562b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -13
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
- # Detect License Plates
 
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
- # Recognize Text from Plates
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
- # Verify Plate
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
- # Custom Video Processor
 
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
- # Return processed frame
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
- # Streamlit UI
 
62
  st.title("Real-Time Car Number Plate Recognition")
63
- st.write("Streamlit with WebRTC for camera streaming.")
64
- webrtc_streamer(key="plate-recognition", video_processor_factory=LicensePlateProcessor)
 
 
 
 
 
 
 
 
 
 
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
+ )