Upload folder using huggingface_hub

__init__.py

@@ -0,0 +1 @@
+# executes when package is imported

api/main.py

@@ -0,0 +1,52 @@
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+import numpy as np
+import cv2
+import base64
+
+from sentinelscan.modeling.infer import CrackModel
+from sentinelscan.utils.viz import overlay_mask
+from sentinelscan.utils.severity import crack_metrics, severity_from_metrics
+
+app = FastAPI(title="SentinelScan API", version="0.1.0")
+
+model = CrackModel(ckpt_path="models/best.pt", size=512)
+
+def _read_image(file_bytes: bytes):
+    arr = np.frombuffer(file_bytes, np.uint8)
+    bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    if bgr is None:
+        raise ValueError("Could not decode image")
+    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+    return rgb
+
+def _to_base64_png(rgb: np.ndarray):
+    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
+    ok, buf = cv2.imencode(".png", bgr)
+    if not ok:
+        raise ValueError("Could not encode image")
+    return base64.b64encode(buf.tobytes()).decode("utf-8")
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    try:
+        rgb = _read_image(await file.read())
+        pred = model.predict(rgb, threshold=0.5)
+
+        m = crack_metrics(pred["mask"])
+        sev = severity_from_metrics(m)
+
+        crack_detected = m["area_px"] > 50  # tiny specks ignored
+
+        overlay = overlay_mask(rgb, pred["mask"])
+        overlay_b64 = _to_base64_png(overlay)
+
+        return JSONResponse({
+            "crack_detected": bool(crack_detected),
+            "confidence": float(pred["confidence"]),
+            "severity": sev if crack_detected else "None",
+            "metrics": m,
+            "overlay_png_base64": overlay_b64,
+        })
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=400)

app/streamlit_app.py

@@ -0,0 +1,40 @@
+import streamlit as st
+import requests
+import base64
+from PIL import Image
+import io
+
+API_URL = "http://localhost:8000/predict"
+
+st.set_page_config(page_title="SentinelScan", layout="centered")
+st.title("🛰️ SentinelScan (Crack Detector v1)")
+
+uploaded = st.file_uploader("Upload an inspection image", type=["jpg","jpeg","png"])
+
+if uploaded:
+    st.subheader("Input")
+    st.image(uploaded, use_container_width=True)
+
+    if st.button("Analyze"):
+        files = {"file": (uploaded.name, uploaded.getvalue(), uploaded.type)}
+        with st.spinner("Running model..."):
+            r = requests.post(API_URL, files=files, timeout=60)
+
+        if r.status_code != 200:
+            st.error(r.text)
+        else:
+            out = r.json()
+            if "error" in out:
+                st.error(out["error"])
+            else:
+                st.subheader("Result")
+                st.write({
+                    "crack_detected": out["crack_detected"],
+                    "severity": out["severity"],
+                    "confidence": out["confidence"],
+                    "metrics": out["metrics"],
+                })
+
+                overlay_bytes = base64.b64decode(out["overlay_png_base64"])
+                overlay_img = Image.open(io.BytesIO(overlay_bytes))
+                st.image(overlay_img, caption="Crack overlay", use_container_width=True)

config.py

@@ -0,0 +1 @@
+# can work with configuration files

data/dataset.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+class CrackSegDataset(Dataset):
+    def __init__(self, images_dir: str, masks_dir: str, transform=None):
+        self.images_dir = Path(images_dir)
+        self.masks_dir = Path(masks_dir)
+        self.transform = transform
+        self.image_paths = sorted([p for p in self.images_dir.glob("*") if p.suffix.lower() in {".jpg",".jpeg",".png"}])
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img_path = self.image_paths[idx]
+        mask_path = self.masks_dir / (img_path.stem + ".png")
+        if not mask_path.exists():
+            raise FileNotFoundError(f"Mask not found for {img_path.name}: {mask_path}")
+
+        image = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+        mask = (mask > 127).astype(np.uint8)  # binarize
+
+        if self.transform is not None:
+            augmented = self.transform(image=image, mask=mask)
+            image, mask = augmented["image"], augmented["mask"]
+
+        # albumentations returns HWC image; convert to CHW float tensor
+        image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
+        mask = torch.from_numpy(mask).unsqueeze(0).float()  # [1,H,W]
+
+        return image, mask

data/transforms.py

@@ -0,0 +1,16 @@
+import albumentations as A
+
+def train_transforms(size=512):
+    return A.Compose([
+        A.Resize(size, size),
+        A.RandomBrightnessContrast(p=0.5),
+        A.GaussianBlur(p=0.2),
+        A.Rotate(limit=15, p=0.4),
+        A.RandomCrop(height=size, width=size, p=0.2),
+        A.GaussNoise(p=0.2),
+    ])
+
+def val_transforms(size=512):
+    return A.Compose([
+        A.Resize(size, size),
+    ])

modeling/infer.py

@@ -0,0 +1,37 @@
+import torch
+import numpy as np
+import cv2
+
+from sentinelscan.modeling.unet import UNet
+
+class CrackModel:
+    def __init__(self, ckpt_path="models/best.pt", device=None, size=512):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.size = size
+        self.model = UNet().to(self.device)
+        ckpt = torch.load(ckpt_path, map_location=self.device)
+        self.model.load_state_dict(ckpt["model_state"])
+        self.model.eval()
+
+    @torch.no_grad()
+    def predict(self, rgb_image: np.ndarray, threshold=0.5):
+        # resize for model
+        img = cv2.resize(rgb_image, (self.size, self.size), interpolation=cv2.INTER_AREA)
+        x = torch.from_numpy(img).permute(2,0,1).float().unsqueeze(0) / 255.0
+        x = x.to(self.device)
+
+        logits = self.model(x)
+        probs = torch.sigmoid(logits).squeeze().cpu().numpy()  # [H,W] float
+        mask = (probs > threshold).astype(np.uint8)
+
+        # confidence: mean prob over predicted crack pixels; fallback to max prob
+        if mask.sum() > 0:
+            conf = float(probs[mask == 1].mean())
+        else:
+            conf = float(probs.max())
+
+        return {
+            "probs": probs,
+            "mask": mask,
+            "confidence": conf,
+        }

modeling/losses.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn.functional as F
+
+def dice_loss(logits, targets, eps=1e-6):
+    probs = torch.sigmoid(logits)
+    num = 2 * (probs * targets).sum(dim=(2,3))
+    den = (probs + targets).sum(dim=(2,3)) + eps
+    dice = num / den
+    return 1 - dice.mean()
+
+def bce_dice_loss(logits, targets, bce_weight=0.5):
+    bce = F.binary_cross_entropy_with_logits(logits, targets)
+    d = dice_loss(logits, targets)
+    return bce_weight * bce + (1 - bce_weight) * d
+
+@torch.no_grad()
+def dice_score(logits, targets, threshold=0.5, eps=1e-6):
+    probs = torch.sigmoid(logits)
+    preds = (probs > threshold).float()
+    num = 2 * (preds * targets).sum(dim=(2,3))
+    den = (preds + targets).sum(dim=(2,3)) + eps
+    return (num / den).mean().item()

modeling/train.py

@@ -0,0 +1,78 @@
+import os
+from dataclasses import dataclass
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from sentinelscan.data.dataset import CrackSegDataset
+from sentinelscan.data.transforms import train_transforms, val_transforms
+from sentinelscan.modeling.unet import UNet
+from sentinelscan.modeling.losses import bce_dice_loss, dice_score
+
+@dataclass
+class TrainConfig:
+    train_images: str = "data/images/train"
+    train_masks: str = "data/masks/train"
+    val_images: str = "data/images/val"
+    val_masks: str = "data/masks/val"
+    out_path: str = "models/best.pt"
+    epochs: int = 25
+    batch_size: int = 8
+    lr: float = 1e-3
+    size: int = 512
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+
+def train(cfg: TrainConfig):
+    os.makedirs(os.path.dirname(cfg.out_path), exist_ok=True)
+
+    train_ds = CrackSegDataset(cfg.train_images, cfg.train_masks, transform=train_transforms(cfg.size))
+    val_ds   = CrackSegDataset(cfg.val_images, cfg.val_masks, transform=val_transforms(cfg.size))
+
+    train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True)
+    val_loader   = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)
+
+    model = UNet().to(cfg.device)
+    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr)
+
+    best_dice = -1.0
+
+    for epoch in range(1, cfg.epochs + 1):
+        model.train()
+        running_loss = 0.0
+
+        for images, masks in tqdm(train_loader, desc=f"Epoch {epoch}/{cfg.epochs} [train]"):
+            images = images.to(cfg.device, non_blocking=True)
+            masks  = masks.to(cfg.device, non_blocking=True)
+
+            opt.zero_grad(set_to_none=True)
+            logits = model(images)
+            loss = bce_dice_loss(logits, masks)
+            loss.backward()
+            opt.step()
+
+            running_loss += loss.item()
+
+        avg_loss = running_loss / max(1, len(train_loader))
+
+        # Validation
+        model.eval()
+        dices = []
+        with torch.no_grad():
+            for images, masks in tqdm(val_loader, desc=f"Epoch {epoch}/{cfg.epochs} [val]"):
+                images = images.to(cfg.device, non_blocking=True)
+                masks  = masks.to(cfg.device, non_blocking=True)
+                logits = model(images)
+                dices.append(dice_score(logits, masks))
+
+        mean_dice = sum(dices) / max(1, len(dices))
+
+        print(f"Epoch {epoch}: loss={avg_loss:.4f} val_dice={mean_dice:.4f}")
+
+        if mean_dice > best_dice:
+            best_dice = mean_dice
+            torch.save({"model_state": model.state_dict(), "cfg": cfg.__dict__}, cfg.out_path)
+            print(f"✅ Saved best model -> {cfg.out_path} (val_dice={best_dice:.4f})")
+
+if __name__ == "__main__":
+    cfg = TrainConfig()
+    train(cfg)

modeling/unet.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def conv_block(in_ch, out_ch):
+    return nn.Sequential(
+        nn.Conv2d(in_ch, out_ch, 3, padding=1),
+        nn.BatchNorm2d(out_ch),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(out_ch, out_ch, 3, padding=1),
+        nn.BatchNorm2d(out_ch),
+        nn.ReLU(inplace=True),
+    )
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1, base=32):
+        super().__init__()
+        self.enc1 = conv_block(in_channels, base)
+        self.enc2 = conv_block(base, base*2)
+        self.enc3 = conv_block(base*2, base*4)
+        self.enc4 = conv_block(base*4, base*8)
+
+        self.pool = nn.MaxPool2d(2)
+
+        self.bottleneck = conv_block(base*8, base*16)
+
+        self.up4 = nn.ConvTranspose2d(base*16, base*8, 2, stride=2)
+        self.dec4 = conv_block(base*16, base*8)
+
+        self.up3 = nn.ConvTranspose2d(base*8, base*4, 2, stride=2)
+        self.dec3 = conv_block(base*8, base*4)
+
+        self.up2 = nn.ConvTranspose2d(base*4, base*2, 2, stride=2)
+        self.dec2 = conv_block(base*4, base*2)
+
+        self.up1 = nn.ConvTranspose2d(base*2, base, 2, stride=2)
+        self.dec1 = conv_block(base*2, base)
+
+        self.head = nn.Conv2d(base, out_channels, 1)
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+
+        b = self.bottleneck(self.pool(e4))
+
+        d4 = self.up4(b)
+        d4 = torch.cat([d4, e4], dim=1)
+        d4 = self.dec4(d4)
+
+        d3 = self.up3(d4)
+        d3 = torch.cat([d3, e3], dim=1)
+        d3 = self.dec3(d3)
+
+        d2 = self.up2(d3)
+        d2 = torch.cat([d2, e2], dim=1)
+        d2 = self.dec2(d2)
+
+        d1 = self.up1(d2)
+        d1 = torch.cat([d1, e1], dim=1)
+        d1 = self.dec1(d1)
+
+        logits = self.head(d1)
+        return logits

utils/severity.py

@@ -0,0 +1,43 @@
+import numpy as np
+import cv2
+from skimage.morphology import skeletonize
+
+def crack_metrics(mask: np.ndarray):
+    # mask: [H,W] 0/1
+    area_px = int(mask.sum())
+
+    # largest connected component
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
+    largest_cc = 0
+    if num_labels > 1:
+        # ignore background label 0
+        largest_cc = int(stats[1:, cv2.CC_STAT_AREA].max())
+
+    # length estimate via skeletonization
+    skel = skeletonize(mask.astype(bool))
+    length_px = int(skel.sum())
+
+    return {
+        "area_px": area_px,
+        "length_px": length_px,
+        "largest_component_px": largest_cc,
+    }
+
+def severity_from_metrics(m):
+    # Tune these thresholds on your validation set
+    area = m["area_px"]
+    length = m["length_px"]
+    largest = m["largest_component_px"]
+
+    score = 0
+    if area > 1500: score += 1
+    if area > 6000: score += 1
+    if length > 600: score += 1
+    if length > 2000: score += 1
+    if largest > 2500: score += 1
+
+    if score >= 4:
+        return "High"
+    if score >= 2:
+        return "Medium"
+    return "Low"

utils/viz.py

@@ -0,0 +1,17 @@
+import numpy as np
+import cv2
+
+def overlay_mask(rgb_image: np.ndarray, mask: np.ndarray, alpha=0.45):
+    # mask: [H,W] 0/1; resize mask to match image
+    h, w = rgb_image.shape[:2]
+    mask_rs = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
+
+    overlay = rgb_image.copy()
+    # red overlay where crack
+    red = np.zeros_like(rgb_image)
+    red[..., 0] = 255
+
+    overlay = np.where(mask_rs[..., None] == 1,
+                       (alpha * red + (1 - alpha) * overlay).astype(np.uint8),
+                       overlay)
+    return overlay