Spaces:
Runtime error
Runtime error
File size: 1,514 Bytes
5794fe9 836176e 778e7ab 836176e 5794fe9 778e7ab 5794fe9 778e7ab 836176e 5794fe9 836176e 5794fe9 836176e 5794fe9 |
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 |
import gradio as gr
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as T
# Define your model class
class UNetClassifier(nn.Module):
def __init__(self, num_classes=1):
super().__init__()
def conv_block(in_c, out_c):
return nn.Sequential(
nn.Conv2d(in_c, out_c, 3, padding=1),
nn.ReLU(),
nn.Conv2d(out_c, out_c, 3, padding=1),
nn.ReLU()
)
self.enc1 = conv_block(3, 64)
self.enc2 = conv_block(64, 128)
self.pool = nn.MaxPool2d(2)
self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(128, num_classes)
def forward(self, x):
x = self.enc1(x)
x = self.pool(x)
x = self.enc2(x)
x = self.pool(x)
x = self.global_pool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return torch.sigmoid(x)
# Load model
model = UNetClassifier()
model.load_state_dict(torch.load("model_weights.pth", map_location="cpu"))
model.eval()
# Image transform
transform = T.Compose([
T.Resize((128, 128)),
T.ToTensor()
])
# Prediction function
def classify_fire(image):
image = image.convert("RGB")
img = transform(image).unsqueeze(0)
with torch.no_grad():
output = model(img)
return "🔥 FIRE" if output.item() > 0.5 else "✅ NO FIRE"
# Gradio UI
gr.Interface(fn=classify_fire, inputs="image", outputs="text", title="Fire Classifier").launch()
|