0801eff

Runtime error

0801eff / app_dict.py

Update app_dict.py

c446819 verified over 2 years ago

1.79 kB

import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import gradio as gr

# 모델 클래스 정의 (ResNet18)
class MyModel(torch.nn.Module):
    def __init__(self, num_classes):
        super(MyModel, self).__init__()
        self.model = models.resnet18(weights=None)  # ResNet18 모델 사용
        self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)  # 마지막 레이어 수정

    def forward(self, x):
        return self.model(x)

# 모델 초기화
num_classes = 6  # 클래스 수에 맞게 설정
model = MyModel(num_classes)

# 모델 가중치 로드 (CPU로 매핑)
model.load_state_dict(torch.load('resnet18_finetuned_2.pth', map_location=torch.device('cpu')), strict=False)
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']

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()]
    if label in ['disposablecup','label'] : result = "Failed"
    else : result = "Successed (+1)"
    # 예측 결과와 원본 출력값을 모두 반환
    return result  # 원본 출력값을 NumPy 배열로 변환하여 반환

# Gradio 인터페이스 설정 (두 개의 출력)
iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs=['label', 'numpy'], live=True)
iface.launch()

	import torch
	import torchvision.transforms as transforms
	from torchvision import models
	from PIL import Image
	import gradio as gr

	# 모델 클래스 정의 (ResNet18)
	class MyModel(torch.nn.Module):
	def __init__(self, num_classes):
	super(MyModel, self).__init__()
	self.model = models.resnet18(weights=None) # ResNet18 모델 사용
	self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes) # 마지막 레이어 수정

	def forward(self, x):
	return self.model(x)

	# 모델 초기화
	num_classes = 6 # 클래스 수에 맞게 설정
	model = MyModel(num_classes)

	# 모델 가중치 로드 (CPU로 매핑)
	model.load_state_dict(torch.load('resnet18_finetuned_2.pth', map_location=torch.device('cpu')), strict=False)
	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']

	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()]
	if label in ['disposablecup','label'] : result = "Failed"
	else : result = "Successed (+1)"
	# 예측 결과와 원본 출력값을 모두 반환
	return result # 원본 출력값을 NumPy 배열로 변환하여 반환

	# Gradio 인터페이스 설정 (두 개의 출력)
	iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs=['label', 'numpy'], live=True)
	iface.launch()