Update app.py

app.py

@@ -16,7 +16,8 @@ transform = transforms.Compose([
 
 # 클래스 이름 리스트 정의
 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():
@@ -25,9 +26,10 @@ def predict(image):
         _, predicted = torch.max(outputs, 1)
     label = class_names[predicted.item()]
     print("Predicted label:", label)
-    
+    if label in class_ok : result = "Mission Successed (+1p)"
+    else : result = "Mission Failed (-)"
     # 예측 결과와 원본 출력값을 모두 반환
-    return label, outputs.numpy()  # 원본 출력값을 NumPy 배열로 변환하여 반환
+    return result, outputs.numpy()  # 원본 출력값을 NumPy 배열로 변환하여 반환
 
 # Gradio 인터페이스 설정 (두 개의 출력)
 iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs=['label', 'numpy'], live=True)