Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,8 +5,8 @@ import numpy as np
 import tempfile
 import torch
 import matplotlib.pyplot as plt
-from transformers import AutoImageProcessor, AutoModelForObjectDetection
 from PIL import Image, ImageDraw
+from transformers import AutoModelForObjectDetection
 
 # Fix cache permission issue
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface'
@@ -15,11 +15,10 @@ model_id = "NaveenKumar5/Solar_panel_fault_detection"
 
 @st.cache_resource
 def load_model():
-    processor = AutoImageProcessor.from_pretrained(model_id)
     model = AutoModelForObjectDetection.from_pretrained(model_id)
-    return processor, model
+    return model
 
-processor, model = load_model()
+model = load_model()
 model.eval()
 
 st.title("🔍 Solar Panel Fault Detection")
@@ -42,19 +41,25 @@ def generate_heatmap(image, boxes):
     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 = processor(images=image, return_tensors="pt")
+        inputs = preprocess_image(image)
 
         with torch.no_grad():
-            outputs = model(**inputs)
+            outputs = model(pixel_values=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)
@@ -82,4 +87,3 @@ if uploaded_file is not None:
 
     elif uploaded_file.type.startswith("video"):
         st.warning("Video support coming soon. For now, please upload an image.")
-