Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,7 +5,7 @@ import numpy as np
 import tempfile
 import torch
 import matplotlib.pyplot as plt
-from transformers import AutoImageProcessor
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
 from PIL import Image, ImageDraw
 
 # Fix cache permission issue
@@ -15,11 +15,11 @@ model_id = "NaveenKumar5/Solar_panel_fault_detection"
 
 @st.cache_resource
 def load_model():
-    extractor = AutoImageProcessor.from_pretrained(model_id)
+    processor = AutoImageProcessor.from_pretrained(model_id)
     model = AutoModelForObjectDetection.from_pretrained(model_id)
-    return extractor, model
+    return processor, model
 
-extractor, model = load_model()
+processor, model = load_model()
 model.eval()
 
 st.title("🔍 Solar Panel Fault Detection")
@@ -45,15 +45,16 @@ def generate_heatmap(image, boxes):
 if uploaded_file is not None:
     if uploaded_file.type.startswith("image"):
         image = Image.open(uploaded_file).convert("RGB")
-        inputs = extractor(images=image, return_tensors="pt")
+        inputs = processor(images=image, return_tensors="pt")
+
         with torch.no_grad():
             outputs = model(**inputs)
 
-        scores = outputs["logits"].softmax(-1)[0].max(-1).values
+        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()
+        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)
@@ -81,3 +82,4 @@ if uploaded_file is not None:
 
     elif uploaded_file.type.startswith("video"):
         st.warning("Video support coming soon. For now, please upload an image.")
+