File size: 1,474 Bytes
e13b4f2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6b13e81
 
e13b4f2
 
 
 
 
 
 
 
5d714eb
 
e13b4f2
a2a9853
e13b4f2
 
5d714eb
e13b4f2
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
import torch
import torchvision.transforms as transforms
from torchvision import models
import gradio as gr

# ์ „์ฒด ๋ชจ๋ธ ๋กœ๋“œ (์ด๋ฏธ ์ •์˜๋œ MyModel ํด๋ž˜์Šค๋Š” ํ•„์š”ํ•˜์ง€ ์•Š์Œ)
model = torch.load('resnet18_finetuned_full.pth')  # ์ „์ฒด ๋ชจ๋ธ์„ ๋ถˆ๋Ÿฌ์˜ด
model.eval()  # ํ‰๊ฐ€ ๋ชจ๋“œ๋กœ ์ „ํ™˜

# ๋ฐ์ดํ„ฐ ๋ณ€ํ™˜ ์ •์˜
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# ํด๋ž˜์Šค ์ด๋ฆ„ ๋ฆฌ์ŠคํŠธ ์ •์˜
class_names = ['disposablecup', 'envmark', 'label', 'mug', 'nonlabel', 'reusablecup']
class_ok = [ 'envmark','mug', 'nonlabel', 'reusablecup']
class_notok = ['disposablecup', 'label']
def predict(image):
    image = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(image)
        print("Raw outputs:", outputs)  # ์›๋ณธ ์ถœ๋ ฅ๊ฐ’ ํ™•์ธ
        _, predicted = torch.max(outputs, 1)
    label = class_names[predicted.item()]
    print("Predicted label:", label)
    if label in class_ok : result = "์ธ์ฆ์ด ์™„๋ฃŒ๋˜์—ˆ์Šต๋‹ˆ๋‹ค. (+1p)"
    else : result = "์ธ์ฆ์— ์‹คํŒจํ–ˆ์Šต๋‹ˆ๋‹ค. (-)"
    # ์˜ˆ์ธก ๊ฒฐ๊ณผ์™€ ์›๋ณธ ์ถœ๋ ฅ๊ฐ’์„ ๋ชจ๋‘ ๋ฐ˜ํ™˜
    return result  # ์›๋ณธ ์ถœ๋ ฅ๊ฐ’์„ NumPy ๋ฐฐ์—ด๋กœ ๋ณ€ํ™˜ํ•˜์—ฌ ๋ฐ˜ํ™˜

# Gradio ์ธํ„ฐํŽ˜์ด์Šค ์„ค์ • (๋‘ ๊ฐœ์˜ ์ถœ๋ ฅ)
iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs=['label'], live=True)
iface.launch()