Spaces:
Sleeping
Sleeping
File size: 2,979 Bytes
bcb590b 1873398 bcb590b 1873398 bcb590b 3d8126c bcb590b 1873398 bcb590b 1873398 3d8126c 1873398 3d8126c bcb590b 1873398 bcb590b 1873398 bcb590b 1873398 bcb590b 1873398 bcb590b 1873398 3d8126c bcb590b 1873398 3d8126c 9f138a6 bcb590b 3d8126c bcb590b 3d8126c bcb590b 3d8126c bcb590b 1873398 bcb590b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
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"
@st.cache_resource
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.")
|