File size: 3,330 Bytes
717af72 | 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 | import gradio as gr
import os
from core.predict import ImageClassifier
from PIL import Image
import cv2
# Set up model path
cwd = os.getcwd()
model_path = os.path.join(cwd, "model", "lenet5_digit_classifier.pth")
# Initialize classifier
classifier = ImageClassifier(model_path=model_path)
# Function to extract and save date digits from cheque image
def save_dates_digits(image_path):
cheque_image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
if cheque_image is None:
raise ValueError("Image could not be read. Check the path or format.")
inverted_image = 255 - cheque_image
cv2.imwrite("inverted_cheque.jpg", inverted_image)
# Coordinates for cropping date region
x, y, w, h = 1790, 100, 460, 80
cropped_date = inverted_image[y:y+h, x:x+w]
# Digit cropping parameters
start_x = 5
start_y = 0
digit_width = 60
digit_height = 70
digit_labels = ['D1', 'D2', 'M1', 'M2', 'Y1', 'Y2', 'Y3', 'Y4']
os.makedirs("./Dates_New", exist_ok=True)
for i, label in enumerate(digit_labels):
digit_img = cropped_date[start_y:start_y + digit_height, start_x + i * digit_width : start_x + (i + 1) * digit_width]
save_path = os.path.join("./Dates_New", f"{label}.jpg")
cv2.imwrite(save_path, digit_img)
# Gradio interface function
def classify_image(image):
image_path = "uploaded_image.jpg"
image.save(image_path)
# Step 1: Extract digit images
save_dates_digits(image_path)
# Step 2: Predict each digit
digit_labels = ['D1', 'D2', 'M1', 'M2', 'Y1', 'Y2', 'Y3', 'Y4']
digit_folder = "./Dates_New"
predicted_digits = []
digit_images = []
for label in digit_labels:
digit_path = os.path.join(digit_folder, f"{label}.jpg")
if os.path.exists(digit_path):
digit_label, _ = classifier.predict(digit_path)
predicted_digits.append(str(digit_label))
digit_images.append((Image.open(digit_path), f"{label}: {digit_label}"))
else:
predicted_digits.append("?")
digit_images.append((None, f"{label}: ?"))
# Step 3: Reconstruct date
day = "".join(predicted_digits[0:2])
month = "".join(predicted_digits[2:4])
year = "".join(predicted_digits[4:8])
reconstructed_date = f"{day}/{month}/{year}"
# Step 4: Load inverted cheque image
inverted_cheque_path = "inverted_cheque.jpg"
inverted_cheque = Image.open(inverted_cheque_path) if os.path.exists(inverted_cheque_path) else None
# Step 5: Return all outputs
return reconstructed_date, image, inverted_cheque, digit_images
# Gradio UI setup
demo = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil", label="Upload Cheque Image"),
outputs=[
gr.Textbox(label="Predicted Date (DD/MM/YYYY)"),
gr.Image(label="Original Image"),
gr.Image(label="Inverted Cheque Image"),
gr.Gallery(label="Digit Predictions", columns=8, height=100)
],
title="Cheque Date Prediction using LeNet-5 CNN",
description="Upload a cheque image to extract and predict handwritten date using a CNN model. Each digit is shown below with its prediction."
)
# Launch the app
if __name__ == "__main__":
demo.launch()
|