alperugurcan's picture
Update app.py
d302107 verified
import streamlit as st
import torch
import torch.nn as nn
import numpy as np
from huggingface_hub import hf_hub_download
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()
)
def forward(self, x):
return self.fc(self.conv(x).view(x.size(0), -1))
@st.cache_resource
def load_model():
model = IcebergClassifier().eval()
model.load_state_dict(torch.load(hf_hub_download("alperugurcan/iceberg","best_iceberg_model.pth"), map_location='cpu'))
return model
st.title('🧊 Simple Ship vs Iceberg Detector')
# Simple numeric inputs
band1 = st.number_input('Enter Band 1 value (-40 to -20)', -40.0, -20.0, -30.0)
band2 = st.number_input('Enter Band 2 value (-35 to -15)', -35.0, -15.0, -25.0)
if st.button('Detect'):
try:
# Create simple 75x75 arrays with the input values
b1 = np.full((75,75), band1)
b2 = np.full((75,75), band2)
# Prepare input tensor
x = torch.FloatTensor(np.stack([b1,b2])).unsqueeze(0)
# Get prediction
model = load_model()
with torch.no_grad():
pred = model(x).item()
# Show result
result = "🧊 ICEBERG" if pred > 0.5 else "🚒 SHIP"
st.success(f"{result} ({pred:.1%})")
except Exception as e:
st.error(f'Error: {str(e)}')