Spaces:
Sleeping
Sleeping
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)) | |
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)}') |