Update app.py

app.py

@@ -1,106 +1,109 @@
 import gradio as gr
 import numpy as np
 from PIL import Image
-import os
+import tensorflow as tf
 
 # ------------------------------
-# 1) 모델 파일 경로 지정
+# 1) 모델 파일 경로
 # ------------------------------
 MODEL_PATH = "crack_detection.h5"
 
-IS_TF = MODEL_PATH.endswith(".h5") or MODEL_PATH.endswith(".keras")
-
 # ------------------------------
 # 2) TensorFlow 모델 로드
 # ------------------------------
-if IS_TF:
-    import tensorflow as tf
-    model = tf.keras.models.load_model(MODEL_PATH)
-    print("🔥 Loaded TensorFlow crack classifier")
-
-# ------------------------------
-# 3) PyTorch 모델 로드
-# ------------------------------
-else:
-    import torch
-    from torch import nn
-    
-    class CNN(nn.Module):
-        # 네가 가진 모델 구조 맞게 조정 필요
-        def __init__(self):
-            super().__init__()
-            self.net = nn.Sequential(
-                nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
-                nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
-                nn.Flatten(),
-                nn.Linear(32 * 56 * 56, 2)  # 이미지 크기 맞게 조정 필요
-            )
-        def forward(self, x):
-            return self.net(x)
+model = tf.keras.models.load_model(MODEL_PATH)
+print("🔥 Loaded TensorFlow crack classifier")
+print("   Input shape :", model.input_shape)
+print("   Output shape:", model.output_shape)
 
-    model = CNN()
-    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
-    model.eval()
-    print("🔥 Loaded PyTorch crack classifier")
+# (None, H, W, C) 형태라고 가정
+input_shape = model.input_shape
+if len(input_shape) != 4:
+    raise ValueError(f"예상치 못한 input_shape: {input_shape}")
+_, H, W, C = input_shape
 
 
 # ------------------------------
-# 4) 전체 예측 함수
+# 3) 예측 함수
+#    항상 JSON을 리턴하도록 try/except
 # ------------------------------
 def predict(img: Image.Image):
-    # 입력 변환
-    img_resized = img.resize((224, 224))
-    arr = np.array(img_resized) / 255.0
+    try:
+        # 1) 이미지 전처리
+        img = img.convert("RGB")
+        img_resized = img.resize((W, H))
+        arr = np.array(img_resized).astype("float32") / 255.0
+        X = np.expand_dims(arr, axis=0)  # (1, H, W, C)
+
+        # 2) 모델 추론
+        raw = model.predict(X)[0]
+        probs = np.array(raw, dtype="float32").flatten()
+
+        # 3) 출력 해석
+        #    - 길이 1  : sigmoid → p_crack
+        #    - 길이 2+ : [p_normal, p_crack] 가정
+        if probs.shape[0] == 1:
+            p_crack = float(probs[0])
+            p_normal = 1.0 - p_crack
 
-    if IS_TF:
-        # TensorFlow
-        X = arr.reshape(1, 224, 224, 3)
-        probs = model.predict(X)[0]  # [p_normal, p_crack]
+            if p_crack >= 0.5:
+                label = "crack"
+                conf = p_crack
+            else:
+                label = "normal"
+                conf = p_normal
 
-    else:
-        # PyTorch
-        import torch
-        X = (
-            torch.tensor(arr)
-            .permute(2, 0, 1)
-            .unsqueeze(0)
-            .float()
-        )
-        probs = torch.softmax(model(X), dim=1).detach().numpy()[0]
+        elif probs.shape[0] >= 2:
+            p_normal = float(probs[0])
+            p_crack = float(probs[1])
 
-    p_normal = float(probs[0])
-    p_crack  = float(probs[1])
+            if p_crack >= p_normal:
+                label = "crack"
+                conf = p_crack
+            else:
+                label = "normal"
+                conf = p_normal
+        else:
+            # 매우 특이한 케이스 → 그냥 argmax
+            idx = int(np.argmax(probs))
+            label = f"class_{idx}"
+            conf = float(probs[idx])
 
-    if p_crack > p_normal:
-        label = "crack"
-        conf  = p_crack
-    else:
-        label = "normal"
-        conf  = p_normal
+        return {
+            "data": [
+                {
+                    "label": label,
+                    "confidence": float(conf),
+                }
+            ]
+        }
 
-    # ------------------------------
-    # 프론트가 요구하는 JSON 구조
-    # ------------------------------
-    return {
-        "data": [
-            {
-                "label": label,
-                "confidence": conf
-            }
-        ]
-    }
+    except Exception as e:
+        # ❗ 여기서 예외를 모두 잡아서 JSON으로 내려줌
+        #    이렇게 해야 HF Space가 500 안 던지고,
+        #    프론트에서 Raw Response/JSON Payload를 볼 수 있음.
+        print("❌ Error in predict():", repr(e))
+        return {
+            "data": [
+                {
+                    "label": "error",
+                    "confidence": 0.0,
+                    "message": str(e),
+                }
+            ]
+        }
 
 
 # ------------------------------
-# 5) Gradio API Interface
+# 4) Gradio Interface
 # ------------------------------
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Image(type="pil", label="Input image"),
     outputs=gr.JSON(label="Detection Result"),
-    title="Crack Detection Classifier",
+    title="Crack Detection Classifier (Keras .h5)",
     description="사진을 업로드하면 균열/정상 여부와 확률(%)을 반환합니다.",
-    flagging_mode="never"
+    flagging_mode="never",
 )
 
 if __name__ == "__main__":