Update app.py

app.py

@@ -1,98 +1,105 @@
 import gradio as gr
-import requests
-from PIL import Image, ImageDraw
+import numpy as np
+from PIL import Image
+import os
+
+# ------------------------------
+# 1) 모델 파일 경로 지정
+# ------------------------------
+MODEL_PATH = "crack_classifier.h5"   # 또는 .pt 로 교체 가능
+
+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 = CNN()
+    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
+    model.eval()
+    print("🔥 Loaded PyTorch crack classifier")
+
+
+# ------------------------------
+# 4) 전체 예측 함수
+# ------------------------------
+def predict(img: Image.Image):
+    # 입력 변환
+    img_resized = img.resize((224, 224))
+    arr = np.array(img_resized) / 255.0
+
+    if IS_TF:
+        # TensorFlow
+        X = arr.reshape(1, 224, 224, 3)
+        probs = model.predict(X)[0]  # [p_normal, p_crack]
 
-# -----------------------------------
-# 1) Roboflow API 설정
-# -----------------------------------
-API_KEY = "UXTLYuI2sw5z7OtCpXdz"
-MODEL_ID = "creacks-eapny/7"
-API_URL = f"https://serverless.roboflow.com/{MODEL_ID}"
-
-
-# -----------------------------------
-# 2) confidence 색상 규칙
-# -----------------------------------
-def get_color(conf):
-    if conf < 0.30:
-        return (0, 255, 0, 120)      # Green
-    elif conf < 0.60:
-        return (255, 165, 0, 120)    # Orange
     else:
-        return (255, 0, 0, 120)      # Red
-
-
-# -----------------------------------
-# 3) Roboflow 호출 + 다중 segmentation 시각화
-# -----------------------------------
-def predict(image: Image.Image):
-    temp_path = "temp.jpg"
-    image.save(temp_path)
-
-    # threshold 낮추기 → crack 감지 활성화
-    with open(temp_path, "rb") as f:
-        response = requests.post(
-            API_URL,
-            files={"file": f},
-            data={
-                "api_key": API_KEY,
-                "confidence": 0.15   # ★ 매우 중요: threshold 낮추기
-            }
+        # PyTorch
+        import torch
+        X = (
+            torch.tensor(arr)
+            .permute(2, 0, 1)
+            .unsqueeze(0)
+            .float()
         )
+        probs = torch.softmax(model(X), dim=1).detach().numpy()[0]
 
-    result = response.json()
-    preds = result.get("predictions", [])
-
-    base = image.convert("RGBA")
-    overlay = Image.new("RGBA", base.size, (0,0,0,0))
-    draw = ImageDraw.Draw(overlay, "RGBA")
+    p_normal = float(probs[0])
+    p_crack  = float(probs[1])
 
-    # 객체 없음 → 정상 처리
-    if not preds:
-        blended = Image.alpha_composite(base, overlay)
-        return blended, {
-            "data": [{
-                "label": "normal",
-                "confidence": 1.0
-            }]
-        }
-
-    # ① 다중 crack 시각화 (for all preds)
-    for p in preds:
-        conf = float(p.get("confidence", 0))
-        color = get_color(conf)
-        points = p.get("points")
-        if points:
-            polygon = [(pt["x"], pt["y"]) for pt in points]
-            draw.polygon(polygon, fill=color)
-
-    blended = Image.alpha_composite(base, overlay)
-
-    # ② 가장 confidence 높은 crack 하나만 프론트로 전달
-    best = max(preds, key=lambda p: p.get("confidence", 0))
-    best_conf = float(best.get("confidence", 0))
-    best_label = best.get("class", "crack")
-
-    return blended, {
-        "data": [{
-            "label": best_label,
-            "confidence": best_conf
-        }]
+    if p_crack > p_normal:
+        label = "crack"
+        conf  = p_crack
+    else:
+        label = "normal"
+        conf  = p_normal
+
+    # ------------------------------
+    # 프론트가 요구하는 JSON 구조
+    # ------------------------------
+    return {
+        "data": [
+            {
+                "label": label,
+                "confidence": conf
+            }
+        ]
     }
 
 
-# -----------------------------------
-# 4) Gradio UI
-# -----------------------------------
+# ------------------------------
+# 5) Gradio API Interface
+# ------------------------------
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Image(label="Crack Heatmap"),
-        gr.JSON(label="Detection Data"),
-    ],
-    title="Crack Detection + Multi-Seg Heatmap",
-    description="여러 균열을 각각 다른 색으로 표시하고, 가장 높은 균열 확률을 UI에 제공합니다.",
+    outputs=gr.JSON(label="Detection Result"),
+    title="Crack Detection Classifier",
+    description="사진을 업로드하면 균열/정상 여부와 확률(%)을 반환합니다.",
     flagging_mode="never"
 )