from flask import Flask, render_template, request
import numpy as np
import base64
import cv2
import os

from ultralytics import YOLO
# model = YOLO('yolov8n.pt')
# model = YOLO('yolo11n.pt')
# development model url
# model = YOLO('adasRoboworld/start_stop_yolo8.pt')
# Deploy model URL
model = YOLO('start_stop_yolo8.pt')
# model = YOLO('runs/detect/train27/weights/last.pt')

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}

app = Flask(__name__)
#this line is for production
application = app

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def predict_on_image(image_stream):
    image = cv2.imdecode( np.asarray(bytearray(image_stream.read()), dtype=np.uint8) , cv2.IMREAD_COLOR)
    
    # //you can pick what class you looking for add classes=0 then it will recognize start
    # results = model.predict(image, classes=0, conf=0.3)
    #recognize all classes that confidence is more then 60%
    results = model.predict(image, conf=0.65)
    print(results[0])
    for image, result in enumerate(results):
        im_bgr = result.plot(conf=True)
        # print(result)
        # print(result.names)

    return im_bgr

@app.route('/', methods=['GET', 'POST'])
def home():
    if request.method == 'POST':
        if 'file' not in request.files:
            return render_template('index.html', error='No file part')

        file = request.files['file']

        if file.filename == '':
            return render_template('index.html', error='No selected file')

        if file and allowed_file(file.filename):

            predicted_image = predict_on_image(file.stream)

            retval, buffer = cv2.imencode('.png', predicted_image)
            detection_img_base64 = base64.b64encode(buffer).decode('utf-8')

            file.stream.seek(0)
            original_img_base64 = base64.b64encode(file.stream.read()).decode('utf-8')

            return render_template('result.html', original_img_data=original_img_base64, detection_img_data=detection_img_base64)

    return render_template('index.html')

# Video stream route
@app.route('/videoStream', methods=['GET', 'POST'])
def homeVideoStream():
    if request.method == 'POST':
        if 'file' not in request.files:
            return render_template('videoStream.html', error='No file part')

        file = request.files['file']

        if file.filename == '':
            return render_template('videoStream.html', error='No selected file')

        if file and allowed_file(file.filename):

            predicted_image = predict_on_image(file.stream)

            retval, buffer = cv2.imencode('.png', predicted_image)
            detection_img_base64 = base64.b64encode(buffer).decode('utf-8')

            file.stream.seek(0)
            original_img_base64 = base64.b64encode(file.stream.read()).decode('utf-8')

            return render_template('resultVideoStream.html', original_img_data=original_img_base64, detection_img_data=detection_img_base64)

    return render_template('videoStream.html')

if __name__ == '__main__':
    os.environ.setdefault('FLASK_ENV', 'development')
    app.run(debug=False, port=5001, host='0.0.0.0')