import numpy as np | |
import time | |
from tensorflow.keras.preprocessing import image | |
# from tensorflow.keras.preprocessing.image import ImageDataGenerator | |
import tensorflow as tf | |
import streamlit as st | |
# with tf.device('/cpu:0'): | |
# Load the saved model | |
model = tf.keras.models.load_model('best_resnet152_model.h5') | |
class_names = {0: '1099_Div', 1: '1099_Int', 2: 'Non_Form', 3: 'w_2', 4: 'w_3'} | |
# print(class_names) | |
# Load and preprocess the image | |
# img_path = '/app/filled_form_1.jpg' | |
def predict(pil_img): | |
# Convert the PIL image to a NumPy array | |
img_array = image.img_to_array(pil_img) | |
img_array = np.expand_dims(img_array, axis=0) | |
img_array /= 255.0 # Rescale pixel values | |
# Predict the class | |
start_time = time.time() | |
predictions = model.predict(img_array) | |
end_time = time.time() | |
predicted_class_index = np.argmax(predictions, axis=1)[0] | |
# Get the predicted class name | |
predicted_class_name = class_names[predicted_class_index] | |
print("Predicted class:", predicted_class_name) | |
# print("Execution time: ", end_time - start_time) | |
return predicted_class_name | |
# import numpy as np | |
# import time | |
# from PIL import Image # Import for PIL image handling | |
# from torchvision import transforms # Import for image preprocessing | |
# import torch | |
# import torch.nn as nn # Import for PyTorch neural networks | |
# import streamlit as st | |
# # Load the PyTorch model (assuming it's saved in PyTorch format) | |
# model = torch.load('./best_resnet152_model.pt') # Replace with your model filename | |
# # Define class names dictionary | |
# class_names = {0: '1099_Div', 1: '1099_Int', 2: 'Non_Form', 3: 'w_2', 4: 'w_3'} | |
# # Define a function for prediction using PyTorch | |
# @st.cache_resource | |
# def predict(pil_img): | |
# # Preprocess the image | |
# preprocess = transforms.Compose([ | |
# transforms.ToTensor(), # Convert to PyTorch tensor | |
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # Normalize based on ImageNet statistics | |
# ]) | |
# img_tensor = preprocess(pil_img) | |
# img_tensor.unsqueeze_(0) # Add batch dimension | |
# # Predict with PyTorch | |
# start_time = time.time() | |
# with torch.no_grad(): # Disable gradient calculation for prediction | |
# predictions = model(img_tensor) | |
# end_time = time.time() | |
# # Get the predicted class | |
# predicted_class_index = torch.argmax(predictions, dim=1).item() | |
# predicted_class_name = class_names[predicted_class_index] | |
# # Print results (optional for debugging) | |
# print("Predicted class:", predicted_class_name) | |
# print("Execution time: ", end_time - start_time) | |
# return predicted_class_name | |