APIMONSTER commited on
Commit
c557d41
Β·
verified Β·
1 Parent(s): 89177eb

Update app.py

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