File size: 4,018 Bytes
4622b44 1279782 4622b44 d885a52 4622b44 9acd233 fa25b8d 4622b44 244c35c fa25b8d f9de43f d885a52 d79e816 d885a52 d79e816 d885a52 4622b44 fa25b8d f9de43f f520e17 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 244c35c 9acd233 244c35c 9acd233 f9de43f 1e12937 f9de43f d79e816 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 124 125 126 127 128 129 130 131 132 133 |
import gradio as gr
import torch
import cv2
import os
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)
model = AutoModel.from_pretrained(
path,
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
# load_in_8bit=True,
low_cpu_mem_usage=True,
use_flash_attn=True,
trust_remote_code=True).eval()
tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
def preprocessing(image):
"""Apply three enhancement filters, including brightness reduction, and resize."""
image = Image.fromarray(np.array(image))
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
width = 800
aspect_ratio = image.height / image.width
height = int(width * aspect_ratio)
image = image.resize((width, height))
return image
def imageRotation(image):
if image.height > image.width:
return image.rotate(90, expand=True)
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,model,tokenizer)
if doc_type == "back":
results = main_b(image,model,tokenizer)
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")
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()
|