import os
from flask import Flask, render_template, request, redirect, url_for,send_from_directory
import cv2
import numpy as np
from transformers import DetrImageProcessor, DetrForObjectDetection
from torchvision.transforms import functional as F
from ultralytics import YOLO
import torch



app = Flask(__name__)
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

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



@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)


@app.route('/', methods=['GET', 'POST'])
def index():
    annotated_image_url = None

    if request.method == 'POST':

        # Load the YOLOv8 model
        yolo_model = YOLO('yolo/yolov8s.pt')

        # Load the DETR model
        processor = DetrImageProcessor.from_pretrained("detr")
        model = DetrForObjectDetection.from_pretrained("detr")

        # Check if a file is selected
        if 'image' not in request.files:
            return redirect(request.url)
        
        image = request.files['image']
        
        # Check if the file has a valid extension
        if image and allowed_file(image.filename):
            constant_filename = 'my_uploaded_image.jpg'  # Specify the constant name
            filename = os.path.join(app.config['UPLOAD_FOLDER'], constant_filename)
            image.save(filename)

            # Load the image for processing
            image = cv2.imread(filename)

            # Perform YOLO object detection and annotation
            yolo_results = yolo_model(image, save=False)
            yolo_image = image.copy()
            yolo_names=yolo_results[0].names
            for row in yolo_results[0].boxes.data:
                x1, y1, x2, y2, score, class_id = row.tolist()
                x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])

                class_name = yolo_names.get(int(class_id), 'Unknown')
                label_text = f"Class: {class_name}, Score: {score:.2f}"
                box_color = (0, 0, 255)
                label_color = (255, 255, 255)

                cv2.rectangle(yolo_image, (x1, y1), (x2, y2), box_color, thickness=2)
                label_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
                label_bottom_left = (x1, y1 - 5)
                label_top_right = (label_bottom_left[0] + label_size[0], label_bottom_left[1] - label_size[1])
                cv2.rectangle(yolo_image, label_bottom_left, label_top_right, box_color, cv2.FILLED)
                cv2.putText(yolo_image, label_text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, label_color, 1, cv2.LINE_AA)




            annotated_filename = 'annotated_my_uploaded_image.jpg'
            annotated_filepath = os.path.join(app.config['UPLOAD_FOLDER'], annotated_filename)
            cv2.imwrite(annotated_filepath, yolo_image)
            annotated_image_url = url_for('uploaded_file', filename=annotated_filename)




            
# Process the image using the processor
            inputs = processor(images=image, return_tensors="pt")
            outputs = model(**inputs)

            # Convert outputs (bounding boxes and class logits) to COCO API format
            # Let's only keep detections with score > 0.9
            target_sizes = torch.tensor([image.shape[:2:]])
            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.3)[0]

            # Convert PIL image to NumPy array for OpenCV
            #image_np = np.array(image)
            #image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
            image_cv2 = image.copy()

            # Define the font for labels
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5
            font_thickness = 1
            font_color = (255, 255, 255)  # White color

            # Iterate over the results and draw bounding boxes and labels using OpenCV
            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
                box = [round(i, 2) for i in box.tolist()]

                # Draw the bounding box
                box = [int(b) for b in box]  # Convert to integers for drawing
                cv2.rectangle(image_cv2, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)  # Red rectangle

                # Draw the label
                label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}"
                label_size = cv2.getTextSize(label_text, font, font_scale, font_thickness)[0]
                label_bottom_left = (box[0], box[1] - 5)  # Adjust label position
                label_top_right = (label_bottom_left[0] + label_size[0], label_bottom_left[1] - label_size[1])
                cv2.rectangle(image_cv2, label_bottom_left, label_top_right, (0, 0, 255), cv2.FILLED)  # Red filled rectangle
                cv2.putText(image_cv2, label_text, (box[0], box[1] - 5), font, font_scale, font_color, font_thickness, cv2.LINE_AA)


            annotated_filename = 'dert_annotated_my_uploaded_image.jpg'
            annotated_filepath = os.path.join(app.config['UPLOAD_FOLDER'], annotated_filename)
            cv2.imwrite(annotated_filepath, image_cv2)
            dertannotated_image_url = url_for('uploaded_file', filename=annotated_filename)





            return render_template('index.html', image1=annotated_image_url ,image2= dertannotated_image_url)









        
    return render_template('index.html', image1=annotated_image_url,image2=annotated_image_url)



if __name__ == '__main__':
    app.run(debug=True,port=7860)