APIMONSTER commited on
Commit
10b8a7d
Β·
verified Β·
1 Parent(s): 28f6af0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +55 -52
app.py CHANGED
@@ -1,7 +1,6 @@
1
  # app.py
2
  import cv2, json, tempfile, re
3
  import numpy as np
4
- np.int = int # work around numpy 2.0 deprecations
5
  import gradio as gr
6
  from ultralytics import YOLO
7
  from paddleocr import PaddleOCR
@@ -9,120 +8,124 @@ from paddleocr import PaddleOCR
9
  # ─── 1) Load models ───────────────────────────────────────────────
10
  yolo = YOLO("models/best.pt")
11
  ocr = PaddleOCR(
12
- det=False, # no built-in detection
13
- rec=True, # recognition only
14
- rec_model_dir="models/ocr_model",
15
- cls=False, # no angle classification
16
  )
17
 
18
  # ─── 2) Turkish plate formatter ────────────────────────────────────
19
  def format_turkish_plate(s: str) -> str:
20
- m = re.match(r"^(\d{2})([A-Z]{1,3})(\d{2,4})$", s.replace(" ",""))
 
21
  return f"{m[1]} {m[2]} {m[3]}" if m else "Unknown"
22
 
23
  # ─── 3) Single‐image pipeline ─────────────────────────────────────
24
  def run_image(img, conf=0.25):
 
25
  bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
26
  res = yolo(bgr, conf=conf)[0]
27
  out = bgr.copy()
28
 
 
29
  for box, yolo_score in zip(res.boxes.xyxy.cpu().numpy(),
30
  res.boxes.conf.cpu().numpy()):
31
- x1,y1,x2,y2 = box.astype(int)
32
  roi = out[y1:y2, x1:x2]
33
- if roi.size == 0:
34
  continue
35
 
36
- plate_img = cv2.resize(roi, (128,32))
 
37
  ocr_result = ocr.ocr(plate_img, cls=True)
38
 
39
- if not ocr_result or not ocr_result[0]:
40
- detected_text = ""
41
- ocr_score = 0.0
42
  else:
43
- detected_text = ocr_result[0][0]
44
- ocr_score = ocr_result[0][1]
45
 
46
- label_text = format_turkish_plate(detected_text)
 
47
  label = f"{label_text} ({ocr_score:.2f})"
 
 
 
48
 
49
- cv2.rectangle(out, (x1,y1), (x2,y2), (0,255,0), 2)
50
- cv2.putText(out, label, (x1, y1-5),
51
- cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
52
-
53
- return cv2.cvtColor(out, cv2.COLOR_BGR2RGB), f"{len(res.boxes)} plate(s)"
54
-
55
 
56
  # ─── 4) Video pipeline ───────────────────────────────────────────
57
  def run_video(video_file, conf=0.25):
58
- cap = cv2.VideoCapture(video_file)
59
- fps = cap.get(cv2.CAP_PROP_FPS)
60
- w, h = int(cap.get(3)), int(cap.get(4))
61
  outfp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
62
- writer= cv2.VideoWriter(outfp, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w,h))
63
 
64
  records = []
65
- idx = 0
66
 
67
  while True:
68
  ret, frame = cap.read()
69
- if not ret:
70
  break
71
- idx += 1
72
- t = idx / fps
73
 
74
  res = yolo(frame, conf=conf)[0]
75
- for (x1,y1,x2,y2) in res.boxes.xyxy.cpu().numpy().astype(int):
76
  roi = frame[y1:y2, x1:x2]
77
- if roi.size == 0:
78
  continue
79
 
80
- plate_img = cv2.resize(roi, (128,32))
81
- rec = ocr.ocr(plate_img, cls=True)[0]
82
-
83
- if rec:
84
- text, ocr_score = rec
85
  else:
86
- text, ocr_score = "", 0.0
87
 
88
- label_text = format_turkish_plate(text)
89
  if label_text != "Unknown":
90
  records.append({
91
- "time_s": round(t,2),
92
  "plate": label_text,
93
- "conf": round(ocr_score,3)
94
  })
95
 
96
- cv2.rectangle(frame, (x1,y1), (x2,y2), (0,255,0), 2)
97
- cv2.putText(frame, label_text, (x1,y1-5),
98
- cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
99
 
100
  writer.write(frame)
101
 
102
  cap.release()
103
  writer.release()
104
 
105
- with open("output.json","w") as f:
 
106
  json.dump(records, f, indent=2)
107
  return outfp
108
 
109
- # ─── 5) Gradio UI ─────────────────────────────────────────────────
110
  with gr.Blocks() as demo:
111
  gr.Markdown("## πŸš— License Plate Detection + Recognition")
 
112
  with gr.Row():
113
  with gr.Column():
114
- img_in = gr.Image(type="numpy", label="Upload Image")
115
- vid_in = gr.File( label="Upload Video (.mp4)")
116
- slider = gr.Slider(0,1,0.25,0.01, label="YOLO Confidence")
117
- btn1 = gr.Button("Run Image")
118
- btn2 = gr.Button("Run Video")
119
  with gr.Column():
120
  img_out = gr.Image(type="numpy", label="Annotated Image")
121
  vid_out = gr.Video(label="Annotated Video")
122
  txt_out = gr.Textbox(label="Status / JSON Path")
123
 
124
- btn1.click(run_image, [img_in, slider], [img_out, txt_out])
125
- btn2.click(run_video, [vid_in, slider], [vid_out, txt_out])
126
 
127
  if __name__ == "__main__":
128
  demo.launch()
 
1
  # app.py
2
  import cv2, json, tempfile, re
3
  import numpy as np
 
4
  import gradio as gr
5
  from ultralytics import YOLO
6
  from paddleocr import PaddleOCR
 
8
  # ─── 1) Load models ───────────────────────────────────────────────
9
  yolo = YOLO("models/best.pt")
10
  ocr = PaddleOCR(
11
+ det=False, # turn off detection
12
+ rec=True, # recognition only
13
+ text_recognition_model_dir="models/ocr_model",
14
+ cls=False # no angle classification
15
  )
16
 
17
  # ─── 2) Turkish plate formatter ────────────────────────────────────
18
  def format_turkish_plate(s: str) -> str:
19
+ # Matches e.g. "06NE944" β†’ "06 NE 944"
20
+ m = re.match(r"^(\d{2})([A-Z]{1,3})(\d{2,4})$", s.replace(" ", ""))
21
  return f"{m[1]} {m[2]} {m[3]}" if m else "Unknown"
22
 
23
  # ─── 3) Single‐image pipeline ─────────────────────────────────────
24
  def run_image(img, conf=0.25):
25
+ # Convert RGB→BGR for YOLO
26
  bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
27
  res = yolo(bgr, conf=conf)[0]
28
  out = bgr.copy()
29
 
30
+ # Loop over each detected box
31
  for box, yolo_score in zip(res.boxes.xyxy.cpu().numpy(),
32
  res.boxes.conf.cpu().numpy()):
33
+ x1, y1, x2, y2 = box.astype(int)
34
  roi = out[y1:y2, x1:x2]
35
+ if roi.size == 0:
36
  continue
37
 
38
+ # Resize to your OCR input shape
39
+ plate_img = cv2.resize(roi, (128, 32))
40
  ocr_result = ocr.ocr(plate_img, cls=True)
41
 
42
+ # Safely unpack [ [poly], (text,score) ] or yield empty
43
+ if ocr_result and ocr_result[0]:
44
+ _, (raw_text, ocr_score) = ocr_result[0]
45
  else:
46
+ raw_text, ocr_score = "", 0.0
 
47
 
48
+ # Format & draw
49
+ label_text = format_turkish_plate(raw_text)
50
  label = f"{label_text} ({ocr_score:.2f})"
51
+ cv2.rectangle(out, (x1, y1), (x2, y2), (0, 255, 0), 2)
52
+ cv2.putText(out, label, (x1, y1 - 5),
53
+ cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
54
 
55
+ # Convert BGR→RGB for display
56
+ return cv2.cvtColor(out, cv2.COLOR_BGR2RGB), f"{len(res.boxes)} plate(s) detected"
 
 
 
 
57
 
58
  # ─── 4) Video pipeline ───────────────────────────────────────────
59
  def run_video(video_file, conf=0.25):
60
+ cap = cv2.VideoCapture(video_file)
61
+ fps = cap.get(cv2.CAP_PROP_FPS)
62
+ w, h = int(cap.get(3)), int(cap.get(4))
63
  outfp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
64
+ writer = cv2.VideoWriter(outfp, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
65
 
66
  records = []
67
+ frame_i = 0
68
 
69
  while True:
70
  ret, frame = cap.read()
71
+ if not ret:
72
  break
73
+ frame_i += 1
74
+ t = frame_i / fps
75
 
76
  res = yolo(frame, conf=conf)[0]
77
+ for (x1, y1, x2, y2) in res.boxes.xyxy.cpu().numpy().astype(int):
78
  roi = frame[y1:y2, x1:x2]
79
+ if roi.size == 0:
80
  continue
81
 
82
+ plate_img = cv2.resize(roi, (128, 32))
83
+ ocr_result = ocr.ocr(plate_img, cls=True)
84
+ if ocr_result and ocr_result[0]:
85
+ _, (raw_text, ocr_score) = ocr_result[0]
 
86
  else:
87
+ raw_text, ocr_score = "", 0.0
88
 
89
+ label_text = format_turkish_plate(raw_text)
90
  if label_text != "Unknown":
91
  records.append({
92
+ "time_s": round(t, 2),
93
  "plate": label_text,
94
+ "conf": round(ocr_score, 3)
95
  })
96
 
97
+ cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
98
+ cv2.putText(frame, label_text, (x1, y1 - 5),
99
+ cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
100
 
101
  writer.write(frame)
102
 
103
  cap.release()
104
  writer.release()
105
 
106
+ # Dump JSON timeline
107
+ with open("output.json", "w") as f:
108
  json.dump(records, f, indent=2)
109
  return outfp
110
 
111
+ # ─── 5) Gradio Web UI ──────────────────────────────────────────────
112
  with gr.Blocks() as demo:
113
  gr.Markdown("## πŸš— License Plate Detection + Recognition")
114
+
115
  with gr.Row():
116
  with gr.Column():
117
+ img_in = gr.Image(type="numpy", label="Upload Image")
118
+ vid_in = gr.File(label="Upload Video (.mp4)")
119
+ slider = gr.Slider(0, 1, 0.25, 0.01, label="YOLO Confidence")
120
+ btn_img = gr.Button("Run Image")
121
+ btn_vid = gr.Button("Run Video")
122
  with gr.Column():
123
  img_out = gr.Image(type="numpy", label="Annotated Image")
124
  vid_out = gr.Video(label="Annotated Video")
125
  txt_out = gr.Textbox(label="Status / JSON Path")
126
 
127
+ btn_img.click(run_image, [img_in, slider], [img_out, txt_out])
128
+ btn_vid.click(run_video, [vid_in, slider], [vid_out, txt_out])
129
 
130
  if __name__ == "__main__":
131
  demo.launch()