import streamlit as st
import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image
import numpy as np
import cv2

# Set page configuration
st.set_page_config(page_title="Solar Panel Fault Detection", layout="wide")

# Title and description
st.title("Solar Panel Fault Detection PoC")
st.write("Upload a thermal image of a solar panel to detect thermal, dust, and power generation faults.")

# Load model and processor
@st.cache_resource
def load_model():
    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
    return processor, model

processor, model = load_model()

# Function to process image and detect faults
def detect_faults(image):
    # Convert PIL image to numpy array
    img_np = np.array(image)
    
    # Convert to RGB if necessary
    if img_np.shape[-1] == 4:
        img_np = img_np[:, :, :3]
    
    # Prepare image for model
    inputs = processor(images=img_np, return_tensors="pt")
    
    # Run inference
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Post-process outputs
    target_sizes = torch.tensor([img_np.shape[:2]])
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
    
    # Initialize fault detection
    faults = {"Thermal Fault": False, "Dust Fault": False, "Power Generation Fault": False}
    annotated_img = img_np.copy()
    
    # Analyze thermal image for faults
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        box = [int(i) for i in box.tolist()]
        # Simulate fault detection based on bounding box and pixel intensity
        roi = img_np[box[1]:box[3], box[0]:box[2]]
        mean_intensity = np.mean(roi)
        
        # Thermal Fault: High intensity (hotspot)
        if mean_intensity > 200:  # Adjust threshold based on thermal image scale
            faults["Thermal Fault"] = True
            cv2.rectangle(annotated_img, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 2)
            cv2.putText(annotated_img, "Thermal Fault", (box[0], box[1]-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
        
        # Dust Fault: Low intensity or irregular patterns
        elif mean_intensity < 100:  # Adjust threshold
            faults["Dust Fault"] = True
            cv2.rectangle(annotated_img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
            cv2.putText(annotated_img, "Dust Fault", (box[0], box[1]-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # Power Generation Fault: Any detected anomaly may indicate reduced efficiency
        if faults["Thermal Fault"] or faults["Dust Fault"]:
            faults["Power Generation Fault"] = True
    
    return annotated_img, faults

# File uploader
uploaded_file = st.file_uploader("Upload a thermal image", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    # Load and display uploaded image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Thermal Image", use_column_width=True)
    
    # Process image
    with st.spinner("Analyzing image..."):
        annotated_img, faults = detect_faults(image)
    
    # Display results
    st.subheader("Fault Detection Results")
    st.image(annotated_img, caption="Annotated Image with Detected Faults", use_column_width=True)
    
    # Show fault summary
    st.write("**Detected Faults:**")
    for fault, detected in faults.items():
        status = "Detected" if detected else "Not Detected"
        color = "red" if detected else "green"
        st.markdown(f"- **{fault}**: <span style='color:{color}'>{status}</span>", unsafe_allow_html=True)
    
    # Provide recommendations
    if any(faults.values()):
        st.subheader("Recommendations")
        if faults["Thermal Fault"]:
            st.write("- **Thermal Fault**: Inspect for damaged components or overheating issues.")
        if faults["Dust Fault"]:
            st.write("- **Dust Fault**: Schedule cleaning to remove dust accumulation.")
        if faults["Power Generation Fault"]:
            st.write("- **Power Generation Fault**: Investigate efficiency issues due to detected faults.")
    else:
        st.write("No faults detected. The solar panel appears to be functioning normally.")

# Footer
st.markdown("---")
st.write("Built with Streamlit and Hugging Face Transformers for Solar Panel Fault Detection PoC")