APIMONSTER commited on
Commit
3a43469
·
verified ·
1 Parent(s): 3e6c137

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +52 -87
app.py CHANGED
@@ -8,180 +8,145 @@ import psutil
8
  from ultralytics import YOLO
9
  from paddleocr import PaddleOCR
10
 
11
- # ─── 0) np.int patch for older PaddleOCR calls
12
- np.int = int
13
 
14
- # ─── 1) Plate character set (digits + uppercase + space)
15
  CHAR_LIST = list("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ ")
16
-
17
- # ─── 2) Load YOLOv8 detector
18
  yolo = YOLO("models/best.pt")
19
 
20
- # ─── 3) Init PaddleOCR recognition-only, override ALL params in-code
21
- ocr = PaddleOCR(
22
- det=False, # disable det on plate crops
23
- rec=True, # recognition-only
24
- rec_model_dir="models/ocr_model",
25
- rec_image_shape="3,32,128", # must match your training
26
- cls=True, # angle classifier
27
- use_angle_cls=True,
28
- use_space_char=True
29
- )
30
- # Force our exact char map (no dict file needed)
31
- ocr.text_recognizer.character = CHAR_LIST
32
  def print_mem_usage(tag=""):
33
  mem = psutil.virtual_memory()
34
  print(f"[{tag}] RAM usage: {mem.used / 1024**2:.2f} MB / {mem.total / 1024**2:.2f} MB ({mem.percent}%)")
35
 
36
-
37
- # ─── 4) Normalize & format OCR output
38
  def normalize_ocr(recs):
39
- """
40
- recs might be:
41
- - [] → no read
42
- - [["ABC123", 0.82]] → default det=False
43
- - [["ABC123", 0.82], ...] → (unlikely here)
44
- - [[box,…], ("ABC123",0.82)] → old det=True style
45
- return text:str, score:float
46
- """
47
  if not recs:
48
  return "", 0.0
49
  first = recs[0]
50
- # case: ["TXT",score]
51
- if isinstance(first, (list,tuple)) and len(first)==2 and isinstance(first[0], str):
52
  return first[0], float(first[1])
53
- # case: [<box>, (<txt>,score)] or [<box>, [txt,score]]
54
- if isinstance(first, (list,tuple)) and len(first)==2 and isinstance(first[1], (list,tuple)):
55
  return first[1][0], float(first[1][1])
56
  return "", 0.0
57
 
58
- # Plaka format kontrolü (Türk plakası değilse orijinal yazı korunur)
59
  def format_plate(s: str) -> str:
60
  s = re.sub(r'[^A-Z0-9]', '', s.upper())
61
  m = re.match(r'^(\d{2})([A-Z]{1,3})(\d{2,4})$', s)
62
  return f"{m.group(1)} {m.group(2)} {m.group(3)}" if m else f"RAW: {s}" if s else "Unknown"
63
 
64
- # Perspektif düzeltme (tek sıra plaka kenarı düz olmayan durumlar için)
65
  def correct_perspective(image, box):
66
  x1, y1, x2, y2 = box
67
  h, w = image.shape[:2]
68
- margin = 5 # biraz dışarıdan al
69
-
70
  x1 = max(0, x1 - margin)
71
  y1 = max(0, y1 - margin)
72
  x2 = min(w, x2 + margin)
73
  y2 = min(h, y2 + margin)
74
-
75
  crop = image[y1:y2, x1:x2]
76
- if crop.size == 0: return None
77
-
78
- # OCR input boyutuna perspektif düzeltme
79
- src_pts = np.float32([
80
- [0, 0], [crop.shape[1], 0],
81
- [crop.shape[1], crop.shape[0]], [0, crop.shape[0]]
82
- ])
83
- dst_pts = np.float32([
84
- [0, 0], [128, 0],
85
- [128, 32], [0, 32]
86
- ])
87
  M = cv2.getPerspectiveTransform(src_pts, dst_pts)
88
  warped = cv2.warpPerspective(crop, M, (128, 32))
89
  return warped
90
 
91
- # Güncellenmiş image fonksiyonu
 
 
 
 
 
 
 
 
 
 
 
 
92
  def run_image(img, conf=0.25):
 
93
  bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
94
  res = yolo(bgr, conf=conf)[0]
95
  out = bgr.copy()
96
 
97
  for box in res.boxes.xyxy.cpu().numpy().astype(int):
98
  warped = correct_perspective(out, box)
99
- if warped is None: continue
100
-
101
  try:
102
  recs = ocr.ocr(warped, det=False, cls=True)
103
- # After OCR call:
104
- gc.collect()
105
- print_mem_usage("After OCR")
106
  except:
107
  recs = []
 
 
108
  txt, score = normalize_ocr(recs)
109
  plate = format_plate(txt)
110
  label = f"{plate} ({score:.2f})"
111
-
112
  x1, y1, x2, y2 = box
113
- cv2.rectangle(out, (x1, y1), (x2, y2), (0,255,0), 2)
114
- cv2.putText(out, label, (x1, y1-8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
115
 
116
  return cv2.cvtColor(out, cv2.COLOR_BGR2RGB), f"{len(res.boxes)} plate(s) detected"
117
 
118
- # Güncellenmiş video fonksiyonu
119
  def run_video(video_file, conf=0.25):
 
120
  cap = cv2.VideoCapture(video_file)
121
  fps = cap.get(cv2.CAP_PROP_FPS) or 30
122
  w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
123
  out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
124
- writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w,h))
125
  records, idx = [], 0
126
 
127
  while True:
128
  ret, frame = cap.read()
129
- if not ret: break
130
- idx += 1; t = idx/fps
 
 
131
  res = yolo(frame, conf=conf)[0]
132
 
133
  for box in res.boxes.xyxy.cpu().numpy().astype(int):
134
  warped = correct_perspective(frame, box)
135
- if warped is None: continue
136
-
137
  try:
138
  recs = ocr.ocr(warped, det=False, cls=True)
139
- # After OCR call:
140
- gc.collect()
141
- print_mem_usage("After OCR")
142
  except:
143
  recs = []
 
 
144
  txt, score = normalize_ocr(recs)
145
  plate = format_plate(txt)
146
- if plate.startswith("RAW:"):
147
- raw_txt = plate[5:]
148
- else:
149
- raw_txt = plate
150
-
151
  if raw_txt != "Unknown":
152
- records.append({"time_s":round(t,2),"plate":raw_txt,"conf":round(score,3)})
153
-
154
  x1, y1, x2, y2 = box
155
- cv2.rectangle(frame, (x1,y1),(x2,y2), (0,255,0), 2)
156
- cv2.putText(frame, plate, (x1, y1-8),
157
- cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
158
 
159
  writer.write(frame)
160
 
161
- cap.release(); writer.release()
162
- with open("output.json","w") as f:
 
163
  json.dump(records, f, indent=2)
164
  return out_path, "Done"
165
 
166
-
167
- # ─── 7) Gradio UI
168
  with gr.Blocks() as demo:
169
  gr.Markdown("## 🚗 License Plate Detection + Recognition")
170
-
171
  with gr.Row():
172
  with gr.Column():
173
  img_in = gr.Image(type="numpy", label="Upload Image")
174
  vid_in = gr.File(label="Upload Video (.mp4)")
175
- conf = gr.Slider(0.0,1.0,0.25,0.01, label="YOLO Confidence")
176
- btn_i = gr.Button("Run Image")
177
- btn_v = gr.Button("Run Video")
178
  with gr.Column():
179
  img_out = gr.Image(type="numpy", label="Annotated Image")
180
  vid_out = gr.Video(label="Annotated Video")
181
- status = gr.Textbox(label="Status / JSON Path")
 
 
 
182
 
183
- btn_i.click(run_image, [img_in,conf], [img_out,status])
184
- btn_v.click(run_video, [vid_in,conf], [vid_out,status])
185
-
186
  if __name__ == "__main__":
187
  demo.launch()
 
8
  from ultralytics import YOLO
9
  from paddleocr import PaddleOCR
10
 
11
+ np.int = int # For backward compatibility
 
12
 
 
13
  CHAR_LIST = list("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ ")
 
 
14
  yolo = YOLO("models/best.pt")
15
 
 
 
 
 
 
 
 
 
 
 
 
 
16
  def print_mem_usage(tag=""):
17
  mem = psutil.virtual_memory()
18
  print(f"[{tag}] RAM usage: {mem.used / 1024**2:.2f} MB / {mem.total / 1024**2:.2f} MB ({mem.percent}%)")
19
 
 
 
20
  def normalize_ocr(recs):
 
 
 
 
 
 
 
 
21
  if not recs:
22
  return "", 0.0
23
  first = recs[0]
24
+ if isinstance(first, (list, tuple)) and len(first) == 2 and isinstance(first[0], str):
 
25
  return first[0], float(first[1])
26
+ if isinstance(first, (list, tuple)) and len(first) == 2 and isinstance(first[1], (list, tuple)):
 
27
  return first[1][0], float(first[1][1])
28
  return "", 0.0
29
 
 
30
  def format_plate(s: str) -> str:
31
  s = re.sub(r'[^A-Z0-9]', '', s.upper())
32
  m = re.match(r'^(\d{2})([A-Z]{1,3})(\d{2,4})$', s)
33
  return f"{m.group(1)} {m.group(2)} {m.group(3)}" if m else f"RAW: {s}" if s else "Unknown"
34
 
 
35
  def correct_perspective(image, box):
36
  x1, y1, x2, y2 = box
37
  h, w = image.shape[:2]
38
+ margin = 5
 
39
  x1 = max(0, x1 - margin)
40
  y1 = max(0, y1 - margin)
41
  x2 = min(w, x2 + margin)
42
  y2 = min(h, y2 + margin)
 
43
  crop = image[y1:y2, x1:x2]
44
+ if crop.size == 0:
45
+ return None
46
+ src_pts = np.float32([[0, 0], [crop.shape[1], 0], [crop.shape[1], crop.shape[0]], [0, crop.shape[0]]])
47
+ dst_pts = np.float32([[0, 0], [128, 0], [128, 32], [0, 32]])
 
 
 
 
 
 
 
48
  M = cv2.getPerspectiveTransform(src_pts, dst_pts)
49
  warped = cv2.warpPerspective(crop, M, (128, 32))
50
  return warped
51
 
52
+ def create_ocr():
53
+ ocr = PaddleOCR(
54
+ det=False,
55
+ rec=True,
56
+ rec_model_dir="models/ocr_model",
57
+ rec_image_shape="3,32,128",
58
+ cls=True,
59
+ use_angle_cls=True,
60
+ use_space_char=True
61
+ )
62
+ ocr.text_recognizer.character = CHAR_LIST
63
+ return ocr
64
+
65
  def run_image(img, conf=0.25):
66
+ ocr = create_ocr()
67
  bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
68
  res = yolo(bgr, conf=conf)[0]
69
  out = bgr.copy()
70
 
71
  for box in res.boxes.xyxy.cpu().numpy().astype(int):
72
  warped = correct_perspective(out, box)
73
+ if warped is None:
74
+ continue
75
  try:
76
  recs = ocr.ocr(warped, det=False, cls=True)
 
 
 
77
  except:
78
  recs = []
79
+ gc.collect()
80
+ print_mem_usage("After OCR")
81
  txt, score = normalize_ocr(recs)
82
  plate = format_plate(txt)
83
  label = f"{plate} ({score:.2f})"
 
84
  x1, y1, x2, y2 = box
85
+ cv2.rectangle(out, (x1, y1), (x2, y2), (0, 255, 0), 2)
86
+ cv2.putText(out, label, (x1, y1 - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
87
 
88
  return cv2.cvtColor(out, cv2.COLOR_BGR2RGB), f"{len(res.boxes)} plate(s) detected"
89
 
 
90
  def run_video(video_file, conf=0.25):
91
+ ocr = create_ocr()
92
  cap = cv2.VideoCapture(video_file)
93
  fps = cap.get(cv2.CAP_PROP_FPS) or 30
94
  w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
95
  out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
96
+ writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
97
  records, idx = [], 0
98
 
99
  while True:
100
  ret, frame = cap.read()
101
+ if not ret:
102
+ break
103
+ idx += 1
104
+ t = idx / fps
105
  res = yolo(frame, conf=conf)[0]
106
 
107
  for box in res.boxes.xyxy.cpu().numpy().astype(int):
108
  warped = correct_perspective(frame, box)
109
+ if warped is None:
110
+ continue
111
  try:
112
  recs = ocr.ocr(warped, det=False, cls=True)
 
 
 
113
  except:
114
  recs = []
115
+ gc.collect()
116
+ print_mem_usage("After OCR")
117
  txt, score = normalize_ocr(recs)
118
  plate = format_plate(txt)
119
+ raw_txt = plate[5:] if plate.startswith("RAW:") else plate
 
 
 
 
120
  if raw_txt != "Unknown":
121
+ records.append({"time_s": round(t, 2), "plate": raw_txt, "conf": round(score, 3)})
 
122
  x1, y1, x2, y2 = box
123
+ cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
124
+ cv2.putText(frame, plate, (x1, y1 - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
 
125
 
126
  writer.write(frame)
127
 
128
+ cap.release()
129
+ writer.release()
130
+ with open("output.json", "w") as f:
131
  json.dump(records, f, indent=2)
132
  return out_path, "Done"
133
 
 
 
134
  with gr.Blocks() as demo:
135
  gr.Markdown("## 🚗 License Plate Detection + Recognition")
 
136
  with gr.Row():
137
  with gr.Column():
138
  img_in = gr.Image(type="numpy", label="Upload Image")
139
  vid_in = gr.File(label="Upload Video (.mp4)")
140
+ conf = gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="YOLO Confidence")
141
+ btn_i = gr.Button("Run Image")
142
+ btn_v = gr.Button("Run Video")
143
  with gr.Column():
144
  img_out = gr.Image(type="numpy", label="Annotated Image")
145
  vid_out = gr.Video(label="Annotated Video")
146
+ status = gr.Textbox(label="Status / JSON Path")
147
+
148
+ btn_i.click(run_image, [img_in, conf], [img_out, status])
149
+ btn_v.click(run_video, [vid_in, conf], [vid_out, status])
150
 
 
 
 
151
  if __name__ == "__main__":
152
  demo.launch()