Captcha / app1.py
Reaumur's picture
Rename app.py to app1.py
551642c verified
raw
history blame
4.54 kB
import streamlit as st
from PIL import Image
import tensorflow as tf
import numpy as np
from keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
import os
# Load custom CTC Layer if necessary
class CTCLayer(tf.keras.layers.Layer):
def __init__(self, name=None):
super().__init__(name=name)
self.loss_fn = tf.keras.backend.ctc_batch_cost
def call(self, y_true, y_pred, input_length, label_length):
# Compute the training-time loss value and add it
# to the layer using `self.add_loss()`.
loss = self.loss_fn(y_true, y_pred, input_length, label_length)
self.add_loss(loss)
# On test time, just return the computed loss
return loss
# Load the trained model with a custom CTC layer if needed
@st.cache_resource
def load_model():
model_path = "model_ocr.h5" # Update with the correct model file path
model = tf.keras.models.load_model(model_path, custom_objects={"CTCLayer": CTCLayer})
return model
model = load_model()
# Menambahkan definisi img_width dan img_height
img_width, img_height = 200, 50 # Ganti sesuai dimensi input gambar yang digunakan oleh model Anda
# Definisikan max_length (misalnya panjang label maksimal)
max_length = 50 # Ganti sesuai dengan panjang label teks maksimal yang diinginkan
# Function to preprocess the image
def prepare_image(img):
# Resize gambar sesuai dengan ukuran yang diharapkan oleh model
img = img.resize((img_width, img_height)) # Resize to (200, 50)
# Konversi gambar ke array
img_array = img_to_array(img)
# Tambahkan dimensi untuk batch (menjadi 1, 50, 200) dan reshape ke bentuk (1, 50, 200, 1)
img_array = np.expand_dims(img_array, axis=0) # Tambahkan dimensi untuk batch
img_array = np.transpose(img_array, (0, 2, 1, 3)) # Mengubah urutan dimensi menjadi (1, 200, 50, 1)
# Menyusun input_length dan label_length untuk model OCR
input_length = np.ones((img_array.shape[0], 1)) * (img_width // 4) # Sesuaikan dengan input panjang
label_length = np.ones((img_array.shape[0], 1)) * max_length # Example label length
# Menambahkan input dummy untuk label (jika perlu untuk prediksi)
dummy_label = np.zeros((img_array.shape[0], max_length)) # Input dummy jika model mengharapkan label input
# Melakukan prediksi
preds = model.predict([img_array, input_length, label_length, dummy_label]) # Berikan 4 input
pred_texts = decode_batch_predictions(preds)
return pred_texts, preds
def decode_batch_predictions(pred):
# Daftar karakter yang digunakan dalam model OCR Anda
characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" # Sesuaikan dengan karakter model Anda
pred_texts = []
for i in range(len(pred)):
# Ambil hasil prediksi untuk sampel i
pred_single = pred[i]
# Ambil argmax untuk setiap langkah (timesteps) untuk mendapatkan indeks karakter
pred_indices = np.argmax(pred_single, axis=-1)
# Ambil hasil prediksi dengan membatasi pada 5 karakter (untuk CAPTCHA dengan panjang 5 karakter)
pred_indices = pred_indices[:5] # Pastikan kita hanya mengambil 5 karakter pertama
# Gabungkan prediksi menjadi string, menghindari padding (-1) dan pastikan panjangnya 5
pred_text = ''.join([characters[int(c)] for c in pred_indices if c != -1])
# Pastikan panjang teks hasil prediksi adalah 5 karakter
while len(pred_text) < 5:
pred_text += " " # Bisa sesuaikan dengan apa yang ingin dilakukan jika panjangnya kurang dari 5
# Append prediksi teks untuk batch
pred_texts.append(pred_text)
return pred_texts
def run():
st.title("OCR Model Deployment")
# Upload image
img_file = st.file_uploader("Choose an Image", type=["jpg", "png"])
if img_file is not None:
img = Image.open(img_file).convert('L') # Convert to grayscale if needed
st.image(img, use_column_width=True)
# Save the uploaded image
upload_dir = './upload_images/'
os.makedirs(upload_dir, exist_ok=True)
save_image_path = os.path.join(upload_dir, img_file.name)
with open(save_image_path, "wb") as f:
f.write(img_file.getbuffer())
# Process the image and make prediction
pred_texts = prepare_image(img)
# Show predicted text
st.success(f"**Predicted Text: {pred_texts[0]}**")
if __name__ == "__main__":
run()