Update app.py

app.py

@@ -1,183 +1,30 @@
 import streamlit as st
-import torch
-import torch.nn as nn
+import torch, torch.nn as nn
 import numpy as np
-import json
-import requests
 from huggingface_hub import hf_hub_download
 
-# Model Architecture
 class IcebergClassifier(nn.Module):
     def __init__(self):
         super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(2, 16, 3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Conv2d(16, 32, 3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Conv2d(32, 64, 3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(2)
-        )
-        self.fc = nn.Sequential(
-            nn.Linear(64 * 9 * 9, 64),
-            nn.ReLU(),
-            nn.Dropout(0.5),
-            nn.Linear(64, 1),
-            nn.Sigmoid()
-        )
+        self.conv = nn.Sequential(nn.Conv2d(2,16,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),
+                                nn.Conv2d(16,32,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),
+                                nn.Conv2d(32,64,3,padding=1),nn.ReLU(),nn.MaxPool2d(2))
+        self.fc = nn.Sequential(nn.Linear(64*9*9,64),nn.ReLU(),nn.Dropout(0.5),
+                              nn.Linear(64,1),nn.Sigmoid())
+    def forward(self, x): return self.fc(self.conv(x).view(x.size(0),-1))
 
-    def forward(self, x):
-        x = self.conv(x)
-        x = x.view(x.size(0), -1)
-        return self.fc(x)
-
-# Model loading with HuggingFace Hub
 @st.cache_resource
-def load_model():
-    try:
-        # Download the model from HuggingFace Hub
-        model_path = hf_hub_download(
-            repo_id="alperugurcan/iceberg",
-            filename="best_iceberg_model.pth"
-        )
-        
-        model = IcebergClassifier()
-        model.load_state_dict(torch.load(model_path, map_location='cpu'))
-        model.eval()
-        return model
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None
-
-def predict(band1, band2):
-    model = load_model()
-    if model is None:
-        return None
-    
+def predict(b1,b2):
+    model = IcebergClassifier().eval()
+    model.load_state_dict(torch.load(hf_hub_download("alperugurcan/iceberg","best_iceberg_model.pth"),map_location='cpu'))
+    return model(torch.FloatTensor(np.stack([np.array(b1).reshape(75,75),np.array(b2).reshape(75,75)])).unsqueeze(0)).item()
+
+st.title('🧊 Iceberg vs Ship Classifier')
+examples = {"Ship":[[-32.5]*75]*75,"Iceberg":[[-28.3]*75]*75,"Strong":[[-35.8]*75]*75}
+for k,v in examples.items(): st.code(f"Example ({k}): {v}")
+b1,b2 = st.text_area('Band 1'),st.text_area('Band 2')
+if st.button('Predict'):
     try:
-        # Transform images to model input format
-        band1 = np.array(band1).reshape(75, 75)
-        band2 = np.array(band2).reshape(75, 75)
-        image = np.stack([band1, band2])
-        image = torch.FloatTensor(image).unsqueeze(0)
-        
-        # Prediction
-        with torch.no_grad():
-            pred = model(image)
-        
-        return pred.item()
-    except Exception as e:
-        st.error(f"Error during prediction: {str(e)}")
-        return None
-
-# Streamlit UI
-def main():
-    st.title('🧊 Iceberg vs Ship Classifier')
-    st.write("""
-    This application uses satellite radar data to predict whether an image shows an iceberg or a ship.
-    The model uses two radar bands (HH and HV polarization) to make its prediction.
-    """)
-
-    # Sidebar with information
-    st.sidebar.header("About")
-    st.sidebar.info("""
-    This model was trained on the Statoil/C-CORE Iceberg Classifier Challenge dataset.
-    It uses radar data from two different polarizations to distinguish between ships and icebergs.
-    """)
-    
-    st.sidebar.header("Input Format")
-    st.sidebar.info("""
-    Each band should be a 75x75 array of radar backscatter values in dB.
-    Values are typically between -50 and 50.
-    """)
-
-    # Main content
-    st.subheader('Radar Image Data')
-    col1, col2 = st.columns(2)
-    
-    with col1:
-        band1_text = st.text_area('Band 1 (HH Polarization)', height=150,
-                                 help='Enter 75x75 array in JSON format')
-    
-    with col2:
-        band2_text = st.text_area('Band 2 (HV Polarization)', height=150,
-                                 help='Enter 75x75 array in JSON format')
-
-    if st.button('🔍 Predict'):
-        if not band1_text or not band2_text:
-            st.warning('Please enter data for both bands')
-            return
-            
-        try:
-            # Parse JSON data
-            band1_data = json.loads(band1_text)
-            band2_data = json.loads(band2_text)
-            
-            # Validate array dimensions
-            if (len(band1_data) != 75 or len(band1_data[0]) != 75 or 
-                len(band2_data) != 75 or len(band2_data[0]) != 75):
-                st.error('Arrays must be 75x75 dimensions')
-                return
-            
-            # Make prediction
-            with st.spinner('Making prediction...'):
-                probability = predict(band1_data, band2_data)
-            
-            if probability is not None:
-                # Show results
-                st.subheader('Prediction Result')
-                
-                # Create columns for the result display
-                result_col1, result_col2 = st.columns(2)
-                
-                with result_col1:
-                    st.metric("Iceberg Probability", f"{probability:.2%}")
-                
-                with result_col2:
-                    if probability > 0.5:
-                        st.success('🧊 ICEBERG')
-                    else:
-                        st.success('🚢 SHIP')
-                
-                # Progress bar
-                st.progress(probability)
-                
-                # Confidence message
-                confidence = abs(probability - 0.5) * 2
-                if confidence > 0.8:
-                    st.write("High confidence prediction")
-                elif confidence > 0.4:
-                    st.write("Medium confidence prediction")
-                else:
-                    st.write("Low confidence prediction")
-                
-        except json.JSONDecodeError:
-            st.error('Please enter valid JSON format data')
-        except Exception as e:
-            st.error(f'An error occurred: {str(e)}')
-
-    # Example usage
-    with st.expander("See example input"):
-        st.code('''
-# Example data format for each band:
-[
-    [-32.5, -31.2, -30.8, ...],  # 75 values
-    [-31.8, -30.9, -31.1, ...],  # 75 values
-    ...                          # 75 rows total
-]
-        ''')
-
-    # Footer
-    st.markdown('---')
-    st.markdown("""
-    <div style='text-align: center'>
-        <p>Made with ❤️ using Streamlit | 
-        <a href='https://huggingface.co/alperugurcan/iceberg'>Model on HuggingFace</a></p>
-    </div>
-    """, unsafe_allow_html=True)
-
-if __name__ == "__main__":
-    main()
+        p = predict(eval(b1),eval(b2))
+        st.write(f"{'🧊 ICEBERG' if p>0.5 else '🚢 SHIP'} ({p:.1%})")
+    except: st.error('Invalid input')