Solar-Panel-Defect-Classification/app.py at main · paras-the-coder/Solar-Panel-Defect-Classification · GitHub

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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
import streamlit as st
from PIL import Image
import numpy as np
import pandas as pd
import os
import requests
import torch
import torchvision.transforms.v2 as transforms

# Import custom PyTorch model and val_transforms from modular src/
from src.model import SolarResNet
from src.dataset import val_transforms

# Page Configuration
st.set_page_config(
    page_title="Solar Panel Defect Classifier (Production)",
    page_icon="☀️",
    layout="wide"
)

CLASS_NAMES = [
    'Bird-drop',
    'Clean',
    'Dusty',
    'Electrical-damage',
    'Physical-Damage',
    'Snow-Covered'
]

# Relative path configuration
MODEL_PATH = 'models/solar_panel_resnet50_best.pt'


# GOOGLE DRIVE CONFIGURATION

# Try to load the file ID from Streamlit secrets (for cloud deployment),
# otherwise fallback to the hardcoded value (for local runs).
try:
    GOOGLE_DRIVE_FILE_ID = st.secrets["GOOGLE_DRIVE_FILE_ID"]
except Exception:
    # Paste your Google Drive File ID here for local running if secrets.toml is not used.
    # Example: GOOGLE_DRIVE_FILE_ID = "1A2B3C4D5E6F"
    GOOGLE_DRIVE_FILE_ID = "YOUR_GOOGLE_DRIVE_FILE_ID_HERE"


def extract_gdrive_id(url_or_id):
    """Extracts the file ID from a full Google Drive URL if provided."""
    if "drive.google.com" in url_or_id:
        if "/d/" in url_or_id:
            return url_or_id.split("/d/")[1].split("/")[0]
        if "id=" in url_or_id:
            return url_or_id.split("id=")[1].split("&")[0]
    return url_or_id

def download_model_from_url(url_or_id, dest_path):
    """Downloads the model from a direct URL (Dropbox/Hugging Face/etc.) or a Google Drive ID."""
    session = requests.Session()
    url_or_id = url_or_id.strip()

    # 1. Google Drive Mode
    if "drive.google.com" in url_or_id or (not url_or_id.startswith("http") and len(url_or_id) > 15):
        file_id = extract_gdrive_id(url_or_id)
        URL = "https://docs.google.com/uc?export=download"

        response = session.get(URL, params={'id': file_id}, stream=True)
        if response.status_code != 200:
            raise ValueError(f"Google Drive returned status code {response.status_code}. "
                             "Please verify that the file sharing permissions are set to 'Anyone with the link can view' (Viewer).")

        token = None
        for key, value in response.cookies.items():
            if key.startswith('download_warning'):
                token = value
                break

        if token:
            params = {'id': file_id, 'confirm': token}
            response = session.get(URL, params=params, stream=True)
            if response.status_code != 200:
                raise ValueError(f"Google Drive confirmation request returned status code {response.status_code}.")

    # 2. Direct Link Mode (Dropbox, Hugging Face, raw link)
    else:
        response = session.get(url_or_id, stream=True)
        if response.status_code != 200:
            raise ValueError(f"Direct download link returned status code {response.status_code}.")

    # Save the file streamingly to prevent memory overflow
    os.makedirs(os.path.dirname(dest_path), exist_ok=True)

    is_html = False
    temp_dest = dest_path + ".tmp"
    with open(temp_dest, "wb") as f:
        for i, chunk in enumerate(response.iter_content(32768)):
            if chunk:
                if i == 0:
                    # Check first chunk for HTML indicators
                    if b"<!DOCTYPE html>" in chunk or b"<html" in chunk:
                        is_html = True
                        break
                f.write(chunk)

    if is_html:
        if os.path.exists(temp_dest):
            os.remove(temp_dest)
        raise ValueError("The download request returned an HTML page instead of the model file. "
                         "If using Google Drive, make sure the file is shared publicly. "
                         "If using Dropbox, make sure the link ends with dl=1.")

    # Rename temp file to actual file path
    if os.path.exists(dest_path):
        os.remove(dest_path)
    os.rename(temp_dest, dest_path)

@st.cache_resource
def load_classifier_model():
    """Loads the trained model, downloading it from Google Drive or a direct URL if missing."""
    # 1. Clean up corrupted or partial downloads (files smaller than 10MB)
    if os.path.exists(MODEL_PATH):
        if os.path.getsize(MODEL_PATH) < 10 * 1024 * 1024:
            st.warning("Detected incomplete or corrupted model weights file. Cleaning up and retrying download...")
            os.remove(MODEL_PATH)

    # 2. Check if model file exists locally
    if not os.path.exists(MODEL_PATH):
        if GOOGLE_DRIVE_FILE_ID == "YOUR_GOOGLE_DRIVE_FILE_ID_HERE":
            st.error("Model file not found! Please configure your `GOOGLE_DRIVE_FILE_ID` in `app.py` or Streamlit Secrets so the app can download it.")
            return None

        try:
            with st.spinner("Downloading model weights... This may take a minute but only happens once."):
                download_model_from_url(GOOGLE_DRIVE_FILE_ID, MODEL_PATH)
            st.success("Model downloaded successfully!")
        except Exception as e:
            st.error("Failed to download model weights. If using Google Drive, make sure permissions are set to 'Anyone with the link'. If using Dropbox, ensure dl=1 is at the end of the URL.")
            st.exception(e)
            # Make sure we clean up any partial file if download crashed
            if os.path.exists(MODEL_PATH):
                os.remove(MODEL_PATH)
            return None

    try:
        # Load the PyTorch model structure
        model = SolarResNet(num_classes=len(CLASS_NAMES))
        # Load weights onto CPU (for safe cross-platform inference)
        model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
        model.eval()
        return model
    except Exception as e:
        st.error(f"Error loading model from relative path '{MODEL_PATH}'. "
                 f"The file might be corrupted. Cleaning up weights so it retries on the next reload.")
        # Auto-delete corrupted file so it retries on next reload
        if os.path.exists(MODEL_PATH):
            os.remove(MODEL_PATH)
        st.exception(e)
        return None


def preprocess_image(image):
    """
    Resizes image to 224x224, converts to PyTorch tensor, and applies
    validation transforms matching the training pipeline.
    """
    # 1. Ensure standard RGB 3-channel
    image = image.convert('RGB')

    # 2. Force resize via PIL bilinear (matching SolarDataset logic)
    image = image.resize((224, 224), Image.Resampling.BILINEAR)

    # 3. Apply standard validation transformations (convert to tensor, float32, normalize)
    tensor_image = val_transforms(image)

    # 4. Add batch dimension -> [1, 3, 224, 224]
    return tensor_image.unsqueeze(0)

def main():
    st.title("☀️ Solar Panel Defect Classification")
    st.markdown("Upload a close-up image of a solar panel (from drone or field inspection) to identify potential defects.")

    # Sidebar configuration
    with st.sidebar:
        st.markdown("### 🖥️ System Status")
        st.markdown("✅ Model Ready  \n✅ Inspection Active")

        st.markdown("---")

        st.markdown("### Detectable Conditions")
        st.markdown("• Bird-Drop  \n• Dusty  \n• Electrical-damage  \n• Physical-Damage  \n• Clean  \n• Snow-Covered")

        st.markdown("---")

        st.markdown("### Tech Stack")
        st.markdown("""
        **AI Model** - ResNet50, PyTorch
        **Processing** - ImageNet Normalization, Data Augmentation
        **Frontend** - Streamlit
        """)

        st.markdown("---")

        st.markdown("### 📊 Performance")
        st.markdown("""
        **Accuracy:** 85%
        **F1 Score:** 85%
        """)

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)

        # Create two columns for clean layout
        col1, col2 = st.columns([1, 1])

        with col1:
            st.image(image, caption='Uploaded Image', use_container_width=True)

        with col2:
            st.subheader("Inspection Results")
            model = load_classifier_model()

            if model:
                with st.spinner('Analyzing solar panel features...'):
                    processed_img = preprocess_image(image).to(torch.device('cpu'))
                    model.eval()
                    with torch.no_grad():
                        outputs = model(processed_img)
                        # Apply softmax to calculate class probabilities
                        probs = torch.softmax(outputs, dim=1)[0].cpu().numpy()

                # Create a DataFrame for visualization
                df = pd.DataFrame({'Class': CLASS_NAMES, 'Probability': probs})
                df = df.sort_values(by='Probability', ascending=False)

                # Display Top Predictions
                st.write("### Predictions Ranking")
                top_3 = df.head(3)

                # Check for critical defects in Top 1 prediction
                top_class = df.iloc[0]['Class']
                top_prob = df.iloc[0]['Probability']

                if top_class == 'Clean':
                    st.success(f"Status: **CLEAN** (Confidence: {top_prob*100:.2f}%)")
                else:
                    st.warning(f"Inspection Status: **{top_class.upper()}** (Confidence: {top_prob*100:.2f}%)")
                    st.info("💡 **Recommendation:** Trigger manual inspection or maintenance for this solar array subset.")

                st.markdown("---")

                for index, row in top_3.iterrows():
                    st.write(f"**{row['Class']}**")
                    st.progress(float(row['Probability']))
                    st.caption(f"{row['Probability']*100:.2f}% Confidence")

                st.divider()

                # Expandable all classes bar chart
                with st.expander("See complete probability distribution"):
                    st.bar_chart(df.set_index('Class'))

                    # Formatting data table
                    formatted_df = df.copy()
                    formatted_df['Probability'] = formatted_df['Probability'].apply(lambda x: f"{x*100:.2f}%")
                    st.table(formatted_df)

if __name__ == "__main__":
    main()