Create app.py

app.py

@@ -0,0 +1,142 @@
+# app.py — Streamlit + WebRTC (Hugging Face Spaces ready)
+import io, numpy as np, torch, torchvision.transforms as T
+from torchvision import models
+from PIL import Image
+import streamlit as st
+import mediapipe as mp
+import cv2
+
+from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
+import av  # needs ffmpeg + pkg-config via packages.txt
+
+st.set_page_config(page_title="Mask Detection (Webcam)", layout="wide")
+st.title("😷 Face Mask Detection — Webcam + Image (HF Spaces)")
+
+LABELS = ["mask", "no_mask"]
+IMG_SIZE = 224
+MEAN = [0.485, 0.456, 0.406]
+STD  = [0.229, 0.224, 0.225]
+
+@st.cache_resource
+def load_model(weights_path: str):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = models.mobilenet_v2(weights=None)
+    model.classifier[1] = torch.nn.Linear(model.last_channel, len(LABELS))
+    state = torch.load(weights_path, map_location="cpu")
+    model.load_state_dict(state, strict=True)
+    model = model.to(device).eval()
+    return model, device
+
+@st.cache_resource
+def get_tf():
+    return T.Compose([
+        T.Resize((IMG_SIZE, IMG_SIZE)),
+        T.ToTensor(),
+        T.Normalize(MEAN, STD),
+    ])
+
+def predict_pil(pil_img, model, device):
+    x = get_tf()(pil_img.convert("RGB")).unsqueeze(0).to(device)
+    with torch.no_grad():
+        probs = torch.softmax(model(x), dim=1)[0].cpu().numpy()
+    i = int(np.argmax(probs))
+    return LABELS[i], float(probs[i]), probs
+
+mp_fd = mp.solutions.face_detection
+@st.cache_resource
+def get_detector():
+    return mp_fd.FaceDetection(model_selection=0, min_detection_confidence=0.5)
+
+def expand_box(x, y, w, h, scale, W, H):
+    cx, cy = x + w/2, y + h/2
+    nw, nh = w*scale, h*scale
+    x1 = int(max(0, cx - nw/2)); y1 = int(max(0, cy - nh/2))
+    x2 = int(min(W, cx + nw/2));  y2 = int(min(H, cy + nh/2))
+    return x1, y1, x2, y2
+
+def annotate_bgr(img_bgr, model, device, conf_thresh=0.6, per_face=True):
+    H, W = img_bgr.shape[:2]
+    out = img_bgr.copy()
+    results = []
+    if not per_face:
+        label, conf, _ = predict_pil(Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)), model, device)
+        color = (0,200,0) if label=="mask" else (0,0,255)
+        cv2.putText(out, f"{label.upper()} {conf:.2f}", (20,60),
+                    cv2.FONT_HERSHEY_SIMPLEX, 1.1, color, 3, cv2.LINE_AA)
+        results.append({"bbox":[0,0,W,H],"label":label,"conf":conf})
+        return out, results
+
+    detector = get_detector()
+    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    det = detector.process(rgb)
+    if not det.detections:
+        return out, results
+
+    for d in det.detections:
+        bb = d.location_data.relative_bounding_box
+        x, y, w, h = int(bb.xmin*W), int(bb.ymin*H), int(bb.width*W), int(bb.height*H)
+        x1, y1, x2, y2 = expand_box(x, y, w, h, 1.25, W, H)
+        crop = img_bgr[max(0,y1):min(H,y2), max(0,x1):min(W,x2)]
+        if crop.size == 0: 
+            continue
+        label, conf, _ = predict_pil(Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)), model, device)
+        if conf < conf_thresh: 
+            continue
+        color = (0,200,0) if label=="mask" else (0,0,255)
+        cv2.rectangle(out, (x1,y1), (x2,y2), color, 2)
+        cv2.putText(out, f"{label.upper()} {conf:.2f}", (x1, max(20,y1-8)),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)
+        results.append({"bbox":[x1,y1,x2,y2], "label":label, "conf":conf})
+    return out, results
+
+def bgr_to_png_bytes(img_bgr):
+    pil = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
+    buf = io.BytesIO(); pil.save(buf, format="PNG"); buf.seek(0); return buf
+
+# Sidebar
+st.sidebar.header("Settings")
+weights_path = st.sidebar.text_input("Model weights (.pt)", value="mask_cls_best.pt")
+conf_thresh = st.sidebar.slider("Confidence threshold", 0.10, 0.99, 0.60, 0.01)
+per_face = st.sidebar.toggle("Per-face boxes (MediaPipe)", value=True)
+
+# Load model
+try:
+    model, device = load_model(weights_path)
+    st.sidebar.success(f"Loaded on {'GPU' if device=='cuda' else 'CPU'}")
+except Exception as e:
+    st.sidebar.error(f"Failed to load weights: {e}")
+    st.stop()
+
+tab1, tab2 = st.tabs(["📷 Image", "🎥 Webcam"])
+
+# Image tab
+with tab1:
+    st.subheader("Image Inference")
+    file = st.file_uploader("Upload an image", type=["jpg","jpeg","png"])
+    if file:
+        pil = Image.open(file).convert("RGB")
+        bgr = cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)
+        out, dets = annotate_bgr(bgr, model, device, conf_thresh=conf_thresh, per_face=per_face)
+        st.image(out, caption="Detections", use_container_width=True)
+        st.download_button("⬇️ Download annotated image", data=bgr_to_png_bytes(out),
+                           file_name="mask_detection.png", mime="image/png")
+
+# Webcam tab (browser camera)
+class FaceMaskTransformer(VideoTransformerBase):
+    def __init__(self):
+        self.model, self.device = model, device
+    def recv(self, frame):
+        img_bgr = frame.to_ndarray(format="bgr24")
+        out, _ = annotate_bgr(img_bgr, self.model, self.device,
+                              conf_thresh=conf_thresh, per_face=per_face)
+        return av.VideoFrame.from_ndarray(out, format="bgr24")
+
+with tab2:
+    st.subheader("Webcam (browser)")
+    st.info("Allow camera access in your browser. If video doesn't appear, open the Space over HTTPS and try Chrome.")
+    webrtc_streamer(
+        key="mask-webrtc",
+        video_transformer_factory=FaceMaskTransformer,
+        media_stream_constraints={"video": True, "audio": False},
+        async_processing=True,
+    )