Update app.py

app.py

@@ -1,70 +1,110 @@
 import gradio as gr
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
 import numpy as np
+import cv2
 from PIL import Image
-import tensorflow as tf
 
-# ------------------------------
-# 1) 모델 파일 경로
-# ------------------------------
-MODEL_PATH = "model.h5"    # GitHub 모델 그대로 사용
+# ---------------------------------------------------
+# 1) 모델 로드
+# ---------------------------------------------------
 
-# ------------------------------
-# 2) TensorFlow 모델 로드
-# ------------------------------
-model = tf.keras.models.load_model(MODEL_PATH)
-print("🔥 Loaded Concrete Crack Classification Model")
-print("   Input shape :", model.input_shape)
-print("   Output shape:", model.output_shape)
+MODEL_PATH = "pretrained_net_G.pth"
 
-# expected input: (None, 227, 227, 3)
-_, H, W, C = model.input_shape
+class DeepCrackNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # 원본 DeepCrack 구조 축약본
+        # pretrained 모델은 UNet-like 구조 기반
+        self.model = torch.load(MODEL_PATH, map_location="cpu")
+        self.model.eval()
+        
+    def forward(self, x):
+        return self.model(x)
 
+net = DeepCrackNet()
+print("🔥 Loaded DeepCrack segmentation model")
 
-# ------------------------------
+# ---------------------------------------------------
+# 2) 이미지 전처리 & 후처리
+# ---------------------------------------------------
+
+transform = T.Compose([
+    T.Resize((256, 256)),
+    T.ToTensor(),
+])
+
+def postprocess_mask(mask_tensor):
+    """tensor → numpy mask (0/255)"""
+    mask = mask_tensor.squeeze().detach().cpu().numpy()
+    mask = (mask * 255).astype(np.uint8)
+    return mask
+
+# ---------------------------------------------------
 # 3) 예측 함수
-# ------------------------------
+# ---------------------------------------------------
+
 def predict(img: Image.Image):
-    # 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, 227, 227, 3)
-
-    # 2) 모델 추론
-    raw = model.predict(X)[0]
-    probs = np.array(raw).flatten()
-
-    # Output structure: [no_crack, crack]
-    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
-
-    # 3) JSON 구조 (프론트 요구사항대로)
-    return {
-        "data": [
-            {
-                "label": label,
-                "confidence": float(conf)
-            }
-        ]
-    }
-
-
-# ------------------------------
-# 4) Gradio API Interface
-# ------------------------------
+    try:
+        # PIL → tensor
+        x = transform(img).unsqueeze(0)
+
+        # DeepCrack forward
+        with torch.no_grad():
+            output = net(x)
+
+        # 모델 출력은 (1, 1, H, W)
+        mask = postprocess_mask(output)
+
+        # 원본 이미지 크기 맞추기
+        orig = np.array(img)
+        mask_resized = cv2.resize(mask, (orig.shape[1], orig.shape[0]))
+
+        # crack 여부 판정
+        crack_pixels = np.sum(mask_resized > 127)
+        total_pixels = mask_resized.size
+        crack_ratio = crack_pixels / total_pixels
+
+        if crack_ratio > 0.01:  # **1% 이상이면 crack**
+            label = "crack"
+            confidence = float(crack_ratio)
+        else:
+            label = "normal"
+            confidence = 1.0 - float(crack_ratio)
+
+        # JSON 출력
+        return {
+            "data": [
+                {
+                    "label": label,
+                    "confidence": confidence
+                }
+            ]
+        }
+
+    except Exception as e:
+        print("❌ ERROR:", e)
+        return {
+            "data": [
+                {
+                    "label": "error",
+                    "confidence": 0.0,
+                    "message": str(e)
+                }
+            ]
+        }
+
+# ---------------------------------------------------
+# 4) Gradio Interface
+# ---------------------------------------------------
+
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.JSON(label="Detection Result"),
-    title="Concrete Crack Classification (227×227 CNN)",
-    description="균열 여부를 판단하고 crack일 때만 확률(%)을 반환합니다.",
+    inputs=gr.Image(type="pil", label="Input"),
+    outputs=gr.JSON(label="Detection"),
+    title="DeepCrack Segmentation API",
+    description="Concrete crack detection using DeepCrack pretrained model.",
     flagging_mode="never",
 )