|
|
import gradio as gr |
|
|
import cv2 |
|
|
import easyocr |
|
|
import numpy as np |
|
|
import os |
|
|
from PIL import Image |
|
|
from ultralytics import YOLO |
|
|
from datetime import datetime |
|
|
|
|
|
|
|
|
model = YOLO("/home/user/app/best.pt") |
|
|
|
|
|
|
|
|
label_map = {0: "Analog", 1: "Digital", 2: "Non-LP"} |
|
|
|
|
|
|
|
|
reader = easyocr.Reader(['bn']) |
|
|
|
|
|
def annotate_frame(frame): |
|
|
input_img = cv2.resize(frame, (640, 640)) |
|
|
results = model(input_img)[0] |
|
|
detections = results.boxes.data.cpu().numpy() |
|
|
|
|
|
for det in detections: |
|
|
if len(det) < 6: |
|
|
continue |
|
|
|
|
|
x1, y1, x2, y2, conf, cls = det |
|
|
x1, y1, x2, y2 = map(int, [x1, y1, x2, y2]) |
|
|
label = label_map.get(int(cls), "Unknown") |
|
|
percent = f"{conf * 100:.2f}%" |
|
|
|
|
|
|
|
|
cv2.rectangle(input_img, (x1, y1), (x2, y2), (255, 0, 0), 2) |
|
|
cv2.putText(input_img, f"{label}: {percent}", (x1, y1 - 10), |
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) |
|
|
|
|
|
|
|
|
cropped = frame[y1:y2, x1:x2] |
|
|
if cropped.size > 0: |
|
|
ocr_result = reader.readtext(cropped) |
|
|
for i, item in enumerate(ocr_result): |
|
|
text = item[1].strip() |
|
|
conf = item[2] |
|
|
cv2.putText(input_img, text, (x1, y2 + 20 + i*25), |
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2) |
|
|
|
|
|
return cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB) |
|
|
|
|
|
def process_input(input_file): |
|
|
file_path = input_file.name |
|
|
ext = os.path.splitext(file_path)[-1].lower() |
|
|
|
|
|
if ext in ['.mp4', '.avi', '.mov']: |
|
|
cap = cv2.VideoCapture(file_path) |
|
|
|
|
|
if not cap.isOpened(): |
|
|
return None, "Could not open video file", "" |
|
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS) |
|
|
width = 640 |
|
|
height = 640 |
|
|
|
|
|
|
|
|
output_path = f"annotated_output.mp4" |
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v') |
|
|
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) |
|
|
|
|
|
while True: |
|
|
ret, frame = cap.read() |
|
|
if not ret: |
|
|
break |
|
|
frame_resized = cv2.resize(frame, (640, 640)) |
|
|
annotated_frame = annotate_frame(frame_resized) |
|
|
bgr_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR) |
|
|
out.write(bgr_frame) |
|
|
|
|
|
cap.release() |
|
|
out.release() |
|
|
|
|
|
return output_path, "Bangla text in video (see frames)", "OCR confidence displayed on frames" |
|
|
|
|
|
else: |
|
|
|
|
|
frame = cv2.imread(file_path) |
|
|
if frame is None: |
|
|
return None, "Invalid image file", "" |
|
|
|
|
|
annotated = annotate_frame(frame) |
|
|
return Image.fromarray(annotated), "Bangla text in image", "" |
|
|
|
|
|
|
|
|
interface = gr.Interface( |
|
|
fn=process_input, |
|
|
inputs=gr.File(type="filepath", label="Upload Image or Video"), |
|
|
outputs=[ |
|
|
gr.Video(label="Output Video"), |
|
|
gr.Image(type="pil", label="Output Image"), |
|
|
gr.Textbox(label="Detected Text (Bangla)"), |
|
|
gr.Textbox(label="Confidence (%)") |
|
|
], |
|
|
title="YOLOv5 License Plate Detector (Bangla OCR)", |
|
|
description="Upload an image or video. Detects license plates and extracts Bangla text using EasyOCR." |
|
|
) |
|
|
|
|
|
interface.launch() |
|
|
|
