import gradio as gr
from transformers import pipeline
from PIL import Image, ImageDraw
import torch

# Initialize the detection pipeline using the DETR architecture
# This model runs locally within the Space environment
try:
    detector = pipeline("object-detection", model="facebook/detr-resnet-50")
except Exception as e:
    detector = None

def analyze_system(image):
    if image is None:
        return None, {"status": "error", "message": "No input signal detected."}
    
    if detector is None:
        return image, {"status": "error", "message": "Model initialization failed."}
    
    # Perform high-precision inference
    predictions = detector(image)
    
    # Prepare drawing context for visual telemetry
    annotated_image = image.copy()
    draw = ImageDraw.Draw(annotated_image)
    
    telemetry_report = []
    
    for pred in predictions:
        box = pred["box"]
        label = pred["label"]
        score = pred["score"]
        
        # Extract spatial coordinates
        xmin, ymin, xmax, ymax = box["xmin"], box["ymin"], box["xmax"], box["ymax"]
        
        # Draw identification borders using a high-contrast industrial green
        draw.rectangle([xmin, ymin, xmax, ymax], outline="#00FF00", width=4)
        
        # Compile telemetry data
        telemetry_report.append({
            "component_class": label,
            "confidence_rating": round(float(score), 4),
            "spatial_coordinates": {
                "xmin": xmin,
                "ymin": ymin,
                "xmax": xmax,
                "ymax": ymax
            }
        })
    
    return annotated_image, telemetry_report

# Construct the Gradio Interface with a technical, utility-focused theme
with gr.Blocks(theme=gr.themes.Monochrome(primary_hue="blue")) as demo:
    gr.Markdown("# 🛰️ Neural Industrial Inspector")
    gr.Markdown("**System Status**: Operational | **Core**: DETR-ResNet-50 Transformer")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(type="pil", label="Optical System Feed")
            run_btn = gr.Button("INITIATE SYSTEM SCAN", variant="primary")
        
        with gr.Column(scale=1):
            output_img = gr.Image(type="pil", label="Visual Diagnostic Overlay")
            output_data = gr.JSON(label="Structured Telemetry Data")
            
    gr.Examples(
        examples=[],
        inputs=input_img
    )
    
    run_btn.click(
        fn=analyze_system, 
        inputs=input_img, 
        outputs=[output_img, output_data]
    )

if __name__ == "__main__":
    demo.launch()