Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,117 +1,83 @@
+import os
 import streamlit as st
-from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
-import torch
-import numpy as np
 import cv2
-from PIL import Image
+import numpy as np
 import tempfile
-import os
+import torch
+import matplotlib.pyplot as plt
+from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
+from PIL import Image, ImageDraw
+
+# Fix cache permission issue
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface'
 
-# Load model once
-@st.cache_resource(show_spinner=True)
+model_id = "NaveenKumar5/Solar_panel_fault_detection"
+
+@st.cache_resource
 def load_model():
-    model_id = "NaveenKumar5/Solar_panel_fault_detection"
     extractor = AutoFeatureExtractor.from_pretrained(model_id)
     model = AutoModelForObjectDetection.from_pretrained(model_id)
-    model.eval()
     return extractor, model
 
 extractor, model = load_model()
+model.eval()
 
-def detect_faults(image: Image.Image):
-    inputs = extractor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    target_sizes = torch.tensor([image.size[::-1]])
-    results = extractor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
-
-    return results
-
-def draw_boxes_and_heatmap(image: Image.Image, results):
-    image_np = np.array(image).copy()
+st.title("🔍 Solar Panel Fault Detection")
+st.write("Upload an image or video to detect faults and view heatmaps.")
 
-    # Create heatmap mask
-    heatmap_mask = np.zeros((image_np.shape[0], image_np.shape[1]), dtype=np.uint8)
+uploaded_file = st.file_uploader("Upload Image or Video", type=["jpg", "png", "mp4", "avi"])
 
-    for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
-        if score < 0.5:
-            continue
-        box = box.int().cpu().numpy()
-        # Draw bounding box
-        cv2.rectangle(image_np, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 2)
-        # Put label text
-        text = f"{label.item()} {score:.2f}"
-        cv2.putText(image_np, text, (box[0], box[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
-        # Fill heatmap mask
-        heatmap_mask[box[1]:box[3], box[0]:box[2]] = 255
+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
 
-    # Create heatmap overlay (apply colormap on mask)
-    heatmap_color = cv2.applyColorMap(heatmap_mask, cv2.COLORMAP_JET)
-    heatmap_color = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB)
+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
 
-    # Overlay heatmap with some transparency
-    overlayed = cv2.addWeighted(image_np, 0.7, heatmap_color, 0.3, 0)
-
-    return Image.fromarray(overlayed)
-
-def process_video(video_path):
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        # Convert frame to PIL
-        pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-        results = detect_faults(pil_frame)
-        frame_out = draw_boxes_and_heatmap(pil_frame, results)
-        # Convert back to BGR for OpenCV video write
-        frame_out = cv2.cvtColor(np.array(frame_out), cv2.COLOR_RGB2BGR)
-        frames.append(frame_out)
-    cap.release()
-    return frames
-
-def save_video(frames, output_path, fps=20):
-    height, width, _ = frames[0].shape
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-    for frame in frames:
-        out.write(frame)
-    out.release()
-
-st.title("Solar Panel Fault Detection with Heatmap")
-
-uploaded_file = st.file_uploader("Upload an image or video", type=["jpg","jpeg","png","mp4","avi"])
-
-if uploaded_file:
+if uploaded_file is not None:
     if uploaded_file.type.startswith("image"):
         image = Image.open(uploaded_file).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-        with st.spinner("Detecting faults..."):
-            results = detect_faults(image)
-            output_image = draw_boxes_and_heatmap(image, results)
-        st.image(output_image, caption="Detection & Heatmap Overlay", use_column_width=True)
-
-        # Show detected faults details
-        st.subheader("Detected Faults")
-        for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
-            if score > 0.5:
-                st.write(f"Fault: **{label.item()}** Confidence: {score:.2f}")
+        inputs = extractor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**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"):
-        tfile = tempfile.NamedTemporaryFile(delete=False)
-        tfile.write(uploaded_file.read())
-        video_path = tfile.name
-
-        st.video(video_path)
-        st.write("Processing video frames. This may take some time depending on video length.")
-
-        with st.spinner("Detecting faults in video..."):
-            frames = process_video(video_path)
-
-            # Save output video
-            output_path = "output.mp4"
-            save_video(frames, output_path)
-
-        st.video(output_path)
-        os.remove(video_path)
+        st.warning("Video support coming soon. For now, please upload an image.")