Update app.py

app.py

@@ -1,111 +1,78 @@
 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
-
-# ---------------------------------------------------
-# 1) 모델 로드
-# ---------------------------------------------------
-
-MODEL_PATH = "pretrained_net_G.pth"
-
-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
-
-# ---------------------------------------------------
+from inference_utils import create_model, inference
+
+# ------------------------------
+# 1) 옵션 클래스 (DeepCrack 기본 설정)
+# ------------------------------
+class Opt:
+    gpu_ids = []
+    isTrain = False
+    checkpoints_dir = "."
+    model = "deepcrack"
+    input_nc = 3
+    output_nc = 3
+    ngf = 64
+    netG = "deepcrack"
+    norm = "batch"
+    use_dropout = False
+    init_type = "normal"
+    init_gain = 0.02
+
+opt = Opt()
+
+# ------------------------------
+# 2) 모델 로드
+# ------------------------------
+print("🔥 Loading DeepCrack model...")
+model = create_model(opt, cp_path="pretrained_net_G.pth")
+print("✅ Model loaded successfully!")
+
+# ------------------------------
 # 3) 예측 함수
-# ---------------------------------------------------
-
+# ------------------------------
 def predict(img: Image.Image):
-    try:
-        # PIL → tensor
-        x = transform(img).unsqueeze(0)
 
-        # DeepCrack forward
-        with torch.no_grad():
-            output = net(x)
+    # PIL → bytes 변환
+    buf = Image.new("RGB", img.size)
+    buf.paste(img)
+    bytes_img = cv2.imencode(".jpg", np.array(buf))[1].tobytes()
 
-        # 모델 출력은 (1, 1, H, W)
-        mask = postprocess_mask(output)
+    # 추론 실행
+    result_img, visuals = inference(
+        model,
+        bytes_img,
+        dim=img.size,
+        unit="px"
+    )
 
-        # 원본 이미지 크기 맞추기
-        orig = np.array(img)
-        mask_resized = cv2.resize(mask, (orig.shape[1], orig.shape[0]))
+    # Pillow로 변환하여 출력
+    out_img = Image.fromarray(result_img)
 
-        # crack 여부 판정
-        crack_pixels = np.sum(mask_resized > 127)
-        total_pixels = mask_resized.size
-        crack_ratio = crack_pixels / total_pixels
+    # JSON은 “균열 여부 + 확률”만 전달
+    prob = float(visuals["fused"].max() / 255.0)
 
-        if crack_ratio > 0.01:  # **1% 이상이면 crack**
-            label = "crack"
-            confidence = float(crack_ratio)
-        else:
-            label = "normal"
-            confidence = 1.0 - float(crack_ratio)
+    has_crack = prob >= 0.5
 
-        # JSON 출력
-        return {
-            "data": [
-                {
-                    "label": label,
-                    "confidence": confidence
-                }
-            ]
-        }
+    return out_img, {
+        "hasCrack": has_crack,
+        "confidence": round(prob, 4)
+    }
 
-    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"),
-    outputs=gr.JSON(label="Detection"),
-    title="DeepCrack Segmentation API",
-    description="Concrete crack detection using DeepCrack pretrained model.",
-    flagging_mode="never",
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Image(label="Crack Segmentation Output"),
+        gr.JSON(label="Prediction")
+    ],
+    title="DeepCrack Segmentation Model",
+    description="Detects crack regions and generates segmentation overlays."
 )
 
 if __name__ == "__main__":