Spaces:
Sleeping
Sleeping
File size: 904 Bytes
25b1bd5 f4cd01c 1b157ea 25b1bd5 |
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 |
import torch
import timm
from PIL import Image
import torchvision.transforms as T
import gradio as gr
# load pretrained backbone and your weights if available
model = timm.create_model("efficientnet_b0", pretrained=True, num_classes=2)
model.eval()
labels = ["No Lipstick", "Lipstick"]
transform = T.Compose([
T.Resize(256),
T.CenterCrop(224),
T.ToTensor(),
T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])
def predict(img):
img = transform(img).unsqueeze(0)
with torch.no_grad():
probs = torch.nn.functional.softmax(model(img), dim=1)[0]
return {labels[i]: float(probs[i]) for i in range(2)}
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=2),
title="π Lipstick Detector",
description="Lightweight EfficientNet-B0 demo without AutoGluon"
)
if __name__ == "__main__":
demo.launch()
|