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()