Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| import tempfile | |
| import torch | |
| import matplotlib.pyplot as plt | |
| from PIL import Image, ImageDraw | |
| from transformers import AutoModelForObjectDetection | |
| # Fix cache permission issue | |
| os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface' | |
| model_id = "NaveenKumar5/Solar_panel_fault_detection" | |
| def load_model(): | |
| model = AutoModelForObjectDetection.from_pretrained(model_id) | |
| return model | |
| model = load_model() | |
| model.eval() | |
| st.title("🔍 Solar Panel Fault Detection") | |
| st.write("Upload an image or video to detect faults and view heatmaps.") | |
| uploaded_file = st.file_uploader("Upload Image or Video", type=["jpg", "png", "mp4", "avi"]) | |
| def draw_boxes(image, boxes, labels, scores): | |
| draw = ImageDraw.Draw(image) | |
| for box, label, score in zip(boxes, labels, scores): | |
| draw.rectangle(box, outline="red", width=2) | |
| draw.text((box[0], box[1] - 10), f"{label}: {score:.2f}", fill="red") | |
| return image | |
| def generate_heatmap(image, boxes): | |
| heatmap = np.zeros((image.shape[0], image.shape[1]), dtype=np.float32) | |
| for box in boxes: | |
| x0, y0, x1, y1 = map(int, box) | |
| heatmap[y0:y1, x0:x1] += 1 | |
| heatmap = np.clip(heatmap / np.max(heatmap), 0, 1) | |
| return heatmap | |
| def preprocess_image(image): | |
| image = image.resize((800, 800)) | |
| image_np = np.array(image).astype(np.float32) / 255.0 | |
| image_tensor = torch.tensor(image_np).permute(2, 0, 1).unsqueeze(0) | |
| return image_tensor | |
| if uploaded_file is not None: | |
| if uploaded_file.type.startswith("image"): | |
| image = Image.open(uploaded_file).convert("RGB") | |
| inputs = preprocess_image(image) | |
| with torch.no_grad(): | |
| outputs = model(pixel_values=inputs) | |
| scores = outputs["logits"].softmax(-1)[0].max(-1).values | |
| keep = scores > 0.5 | |
| boxes = outputs["pred_boxes"][0][keep].cpu().numpy() | |
| labels = outputs["logits"].argmax(-1)[0][keep].cpu().numpy() | |
| scores = scores[keep].cpu().numpy() | |
| image_np = np.array(image) | |
| height, width = image_np.shape[:2] | |
| abs_boxes = [] | |
| for box in boxes: | |
| cx, cy, w, h = box | |
| x0 = int((cx - w / 2) * width) | |
| y0 = int((cy - h / 2) * height) | |
| x1 = int((cx + w / 2) * width) | |
| y1 = int((cy + h / 2) * height) | |
| abs_boxes.append([x0, y0, x1, y1]) | |
| # Draw boxes and labels | |
| boxed_image = draw_boxes(image.copy(), abs_boxes, labels, scores) | |
| st.image(boxed_image, caption="Detected Faults", use_column_width=True) | |
| # Generate and show heatmap | |
| heatmap = generate_heatmap(image_np, abs_boxes) | |
| fig, ax = plt.subplots() | |
| ax.imshow(image_np) | |
| ax.imshow(heatmap, cmap="jet", alpha=0.5) | |
| ax.axis("off") | |
| st.pyplot(fig) | |
| elif uploaded_file.type.startswith("video"): | |
| st.warning("Video support coming soon. For now, please upload an image.") | |