File size: 3,740 Bytes
4622b44 d885a52 4622b44 9acd233 fa25b8d 4622b44 fa25b8d f9de43f d885a52 f9de43f d885a52 4622b44 fa25b8d f9de43f af1f46a fa25b8d f9de43f 4622b44 f9de43f 4622b44 f9de43f 4622b44 f9de43f 4622b44 f9de43f 04dd650 f9de43f 04dd650 f9de43f 4622b44 1e12937 f9de43f d144786 f9de43f 4622b44 1e12937 f9de43f 9acd233 f9de43f 1e12937 f9de43f ce943c6 4622b44 1e12937 4622b44 f9de43f 4622b44 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image, ImageEnhance
from ultralytics import YOLO
from backPrompt import main as main_b
from frontPrompt import main as main_f
model_path = "best.pt"
model = YOLO(model_path)
def preprocessing(image):
"""Apply three enhancement filters, including brightness reduction, and resize."""
image = Image.fromarray(np.array(image))
# Apply enhancements
image = ImageEnhance.Sharpness(image).enhance(2.0) # Increase sharpness
image = ImageEnhance.Contrast(image).enhance(1.5) # Increase contrast
image = ImageEnhance.Brightness(image).enhance(0.8) # Reduce brightness
# Resize image to 800px width while maintaining aspect ratio
width = 800
aspect_ratio = image.height / image.width
height = int(width * aspect_ratio)
image = image.resize((width, height))
return image
def imageRotation(image):
"""Dummy function for image rotation."""
return image
def detect_document(image):
"""Detects front and back of the document using YOLO."""
image = np.array(image)
results = model(image, conf=0.85)
detected_classes = set()
labels = []
bounding_boxes = []
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = box.conf[0]
cls = int(box.cls[0])
class_name = model.names[cls]
detected_classes.add(class_name)
label = f"{class_name} {conf:.2f}"
labels.append(label)
bounding_boxes.append((x1, y1, x2, y2, class_name, conf)) # Store bounding box with class and confidence
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
possible_classes = {"front", "back"}
missing_classes = possible_classes - detected_classes
if missing_classes:
labels.append(f"Missing: {', '.join(missing_classes)}")
return Image.fromarray(image), labels, bounding_boxes
def crop_image(image, bounding_boxes):
"""Crops detected bounding boxes from the image."""
cropped_images = {}
image = np.array(image)
for (x1, y1, x2, y2, class_name, conf) in bounding_boxes:
cropped = image[y1:y2, x1:x2]
cropped_images[class_name] = Image.fromarray(cropped)
return cropped_images
def vision_ai_api(image, doc_type):
if doc_type == "front":
results = main_f(image)
if doc_type == "back":
results = main_b(image)
return results
def predict(image):
"""Pipeline: Preprocess -> Detect -> Crop -> Vision AI API."""
processed_image = preprocessing(image)
rotated_image = imageRotation(processed_image) # Placeholder for rotation
detected_image, labels, bounding_boxes = detect_document(rotated_image)
cropped_images = crop_image(rotated_image, bounding_boxes)
# Call Vision AI separately for front and back if detected
front_result, back_result = None, None
if "front" in cropped_images:
front_result = vision_ai_api(cropped_images["front"], "front")
if "back" in cropped_images:
back_result = vision_ai_api(cropped_images["back"], "back")
# Combine API results into a single response
api_results = {
"front": front_result,
"back": back_result
}
single_image = cropped_images.get("front") or cropped_images.get("back") or detected_image
return single_image, labels, api_results
iface = gr.Interface(
fn=predict,
inputs="image",
outputs=["image", "text", "json"],
title="License Field Detection (Front & Back Card)"
)
iface.launch()
|