import gradio as gr
from roboflow import Roboflow
import os
import tempfile
from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np

# Initialize Roboflow
rf = Roboflow(api_key="E5qhgf3ZimDoTx5OfgZ8")
project = rf.workspace().project("newhassae")

def get_model(version):
    return project.version(version).model

def preprocess_image(img, version):
    # Initial crop for all images
    img = img.crop((682, 345, 682+2703, 345+1403))
    
    # Model specific processing
    if version == 1:
        return img.resize((640, 640))
    elif version == 2:
        return img
    elif version == 3:
        width, height = img.size
        left = (width - 640) // 2
        top = (height - 640) // 2
        right = left + 640
        bottom = top + 640
        return img.crop((left, top, right, bottom))
    return img

def process_images(image_files, version):
    model = get_model(version)
    results = []
    if not isinstance(image_files, list):
        image_files = [image_files]
        
    for image_file in image_files:
        try:
            with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file:
                temp_file.write(image_file)
                temp_path = temp_file.name
            
            img = Image.open(temp_path)
            processed_img = preprocess_image(img, version)
            
            processed_temp = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False)
            processed_img.save(processed_temp.name)
            
            try:
                prediction = model.predict(processed_temp.name).json()
                predicted_class = prediction["predictions"][0]["predictions"][0]["class"]
                confidence = f"{float(prediction['predictions'][0]['predictions'][0]['confidence']) * 100:.1f}%"
            except Exception as e:
                predicted_class = "Error"
                confidence = "N/A"
            
            if processed_img.mode != 'RGB':
                processed_img = processed_img.convert('RGB')
                
            labeled_img = add_label_to_image(processed_img, predicted_class, confidence)
            
            top_result = {
                "predicted_class": predicted_class,
                "confidence": confidence
            }
            
            results.append((labeled_img, top_result))
            
        except Exception as e:
            gr.Warning(f"Error processing image: {str(e)}")
            continue
        finally:
            if 'temp_path' in locals():
                os.unlink(temp_path)
            if 'processed_temp' in locals():
                os.unlink(processed_temp.name)
    
    return results if results else [(Image.new('RGB', (400, 400), 'grey'), {"predicted_class": "Error", "confidence": "N/A"})]



def add_label_to_image(image, prediction, confidence):
    # Convert PIL image to OpenCV format
    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # Image dimensions
    img_height, img_width = img_cv.shape[:2]
    padding = int(img_width * 0.02)
    
    # Rectangle dimensions
    rect_height = int(img_height * 0.15)
    rect_width = img_width - (padding * 2)
    
    # Convert confidence percentage string to float
    confidence_value = float(confidence.strip('%'))
    
    # Set color based on confidence (BGR format)
    rect_color = (0, 255, 0) if confidence_value >= 80 else (0, 0, 255)  # Green if >=80%, Red otherwise
    
    # Draw colored rectangle
    cv2.rectangle(img_cv, 
                 (padding, padding), 
                 (padding + rect_width, padding + rect_height), 
                 rect_color, 
                 -1)
    
    text = f"{prediction}: {confidence}"
    
    # Text settings
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 3.0
    thickness = 8
    
    # Get text size and position
    (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
    text_x = padding + (rect_width - text_width) // 2
    text_y = padding + (rect_height + text_height) // 2
    
    # Draw white text
    cv2.putText(img_cv, text, (text_x, text_y), font, font_scale, (255, 255, 255), thickness)
    
    # Convert back to PIL
    img_pil = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))
    return img_pil

def display_results(image_files, version):
    results = process_images(image_files, version)
    output_images = [res[0] for res in results]
    predictions = [res[1] for res in results]
    
    return output_images, predictions
# Create Gradio interface
with gr.Blocks() as demo:
    gr.HTML("""
        <div style="text-align: center; margin-bottom: 1rem">
            <img src="https://haeab.se/wp-content/uploads/2023/12/ad.png" alt="Logo" style="height: 100px;">
        </div>
    """)
    gr.Markdown("Hans Andersson Entrepenad")
    
    with gr.Row():
        with gr.Column():
            model_version = gr.Slider(
                minimum=1,
                maximum=4,
                step=1,
                value=1,
                label="Model Version",
                interactive=True
            )
            image_input = gr.File(
                label="Upload Image(s)",
                file_count="multiple",
                type="binary"
            )
        
        with gr.Column():
            image_output = gr.Gallery(label="Processed Images")
            text_output = gr.JSON(
                label="Top Predictions",
                height=400,  # Increases height
                container=True,  # Adds a container around the JSON
                    show_label=True
                )
    
    submit_btn = gr.Button("Analyze Images")
    submit_btn.click(
        fn=display_results,
        inputs=[image_input, model_version],
        outputs=[image_output, text_output]
    )
demo.launch(share=True, debug=True, show_error=True)