from flask import Flask, request, jsonify, render_template
from flask_cors import CORS
import cv2
import easyocr
import numpy as np
import tensorflow as tf

app = Flask(__name__)
CORS(app)

# Load your pre-trained model and set up the test set for phone recognition
test_set = tf.keras.utils.image_dataset_from_directory(
    "/app/testing_data",
    labels="inferred",
    label_mode="categorical",
    class_names=None,
    color_mode="rgb",
    batch_size=32,
    image_size=(64, 64),
    shuffle=True,
    seed=None,
    validation_split=None,
    subset=None,
    interpolation="bilinear",
    follow_links=False,
    crop_to_aspect_ratio=False
)

# Load the CNN model for phone recognition
cnn = tf.keras.models.load_model("/app/trained_model (1).h5")

# Instantiate text reader for OCR
picture_read = easyocr.Reader(["en", "ar"], gpu=False)

def get_class_names():
    # Load the class names from the test_set
    return test_set.class_names

class_names = get_class_names()

# Route for rendering the index.html template
@app.route("/")
def index():
    # Pass the class names to the HTML template
    return render_template("index.html", class_names=class_names)

# Route for phone recognition and processing image
@app.route("/process_image", methods=["POST"])
def process_image():
    # Get the uploaded file from the request
    uploaded_file = request.files["file"]

    # Save the file to a temporary location
    temp_file_path = "/app/temp_image.png"
    uploaded_file.save(temp_file_path)

    # Read and preprocess the image for phone model recognition
    img = cv2.imread(temp_file_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (64, 64))

    # Make predictions for phone model
    input_arr = np.array([img])
    predictions = cnn.predict(input_arr)

    # Get the index with the highest probability
    result_index = np.argmax(predictions)

    # Display the predicted class for phone model
    result_phone_model = class_names[result_index]

    # Process image using OCR
    picture_results = picture_read.readtext(temp_file_path)
    results_ocr = []

    conf_threshold = 0.2
    for y in picture_results:
        if y[2] > conf_threshold:
            text = y[1]
            results_ocr.append({"text": text})

    # Return JSON response with the results
    return jsonify({"result_phone_model": result_phone_model, "results_ocr": results_ocr})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000, debug=True)