File size: 2,924 Bytes
1393304 6158c53 1393304 3e477f8 1393304 3e477f8 1393304 3e477f8 1393304 3e477f8 1393304 3e477f8 1393304 3e477f8 1393304 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance
from ultralytics import YOLO
import cv2
# Load YOLO model
model_path = "./best.pt"
modelY = YOLO(model_path)
modelY.to('cpu')
# Preprocessing function
def preprocessing(image):
if image.mode != 'RGB':
image = image.convert('RGB')
image = ImageEnhance.Sharpness(image).enhance(2.0)
image = ImageEnhance.Contrast(image).enhance(1.5)
image = ImageEnhance.Brightness(image).enhance(0.8)
width = 448
aspect_ratio = image.height / image.width
height = int(width * aspect_ratio)
return image.resize((width, height))
# YOLO document detection and cropping
def detect_and_crop_document(image):
image_np = np.array(image)
results = modelY(image_np, conf=0.80, device='cpu')
cropped_images = []
predictions = []
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = int(box.conf[0] * 100) # Convert confidence to percentage
cls = int(box.cls[0])
class_name = modelY.names[cls].capitalize() # Capitalize class names
cropped_image_np = image_np[y1:y2, x1:x2]
cropped_image = Image.fromarray(cropped_image_np)
cropped_images.append(cropped_image)
predictions.append(f"Detected: STNK {class_name} -- (Confidence: {conf}%)")
if not cropped_images:
return None, "No document detected"
return cropped_images, predictions
# Gradio interface
def process_image(image):
preprocessed_image = preprocessing(image)
cropped_images, predictions = detect_and_crop_document(preprocessed_image)
if cropped_images:
return cropped_images, '\n'.join(predictions)
return None, "No document detected"
with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-size: 16px; padding: 10px 20px; border-radius: 8px;}") as demo:
gr.Markdown(
"""
<h1 style="text-align: center; color: #4caf50;">📜 License Registration Classification</h1>
<p style="text-align: center; font-size: 18px;">Upload an image and let the YOLO model detect and crop license documents automatically.</p>
"""
)
with gr.Row():
with gr.Column(scale=1, min_width=300):
input_image = gr.Image(type="pil", label="Upload License Image", interactive=True)
with gr.Row():
clear_btn = gr.Button("Clear")
submit_btn = gr.Button("Detect Document")
with gr.Column(scale=2):
output_image = gr.Gallery(label="Cropped Documents", interactive=False)
output_text = gr.Textbox(label="Detection Result", interactive=False)
submit_btn.click(process_image, inputs=input_image, outputs=[output_image, output_text])
clear_btn.click(lambda: (None, ""), outputs=[output_image, output_text])
demo.launch() |