Spaces:
Sleeping
Sleeping
| 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 | |
| # Load YOLO model | |
| model = YOLO("/home/user/app/best.pt") | |
| # Label map | |
| label_map = {0: "Analog", 1: "Digital", 2: "Non-LP"} | |
| # EasyOCR Bengali | |
| 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}%" | |
| # Draw box and label | |
| 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) | |
| # OCR | |
| 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, None, "Could not open video file", "" | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| frame_skip = 5 | |
| frame_id = 0 | |
| output_path = "annotated_output.mp4" | |
| out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (640, 640)) | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| if frame_id % frame_skip != 0: | |
| frame_id += 1 | |
| continue | |
| frame_id += 1 | |
| frame = cv2.resize(frame, (640, 640)) | |
| annotated = annotate_frame(frame) | |
| out.write(cv2.cvtColor(annotated, cv2.COLOR_RGB2BGR)) | |
| cap.release() | |
| out.release() | |
| return output_path, None, "Bangla text in video (see frames)", "OCR confidence displayed" | |
| else: | |
| frame = cv2.imread(file_path) | |
| if frame is None: | |
| return None, None, "Invalid image", "" | |
| frame = cv2.resize(frame, (640, 640)) | |
| annotated = annotate_frame(frame) | |
| pil_img = Image.fromarray(annotated) | |
| return None, pil_img, "Bangla text in image", "OCR confidence 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() | |