Initial upload of YOLO-style face detection model

MIT License.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 YOUR NAME
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

README.md

@@ -1,3 +1,15 @@
----
-license: mit
----
+
+# Face Detection (WIDER FACE, YOLO-style)
+
+## Train
+```
+python train.py
+```
+
+## Inference
+```
+python inference.py
+```
+
+## Model
+- S=20, B=3 YOLO-style bounding boxes

dataset.py

@@ -0,0 +1,50 @@
+
+import os
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+
+class WiderFaceMultiBoxDataset(Dataset):
+    def __init__(self, images_dir, annotation_file, S=20, B=3, transform=None):
+        self.images_dir = images_dir
+        self.annotation_file = annotation_file
+        self.transform = transform
+        self.S = S
+        self.B = B
+        self.data = self.parse_annotations()
+
+    def parse_annotations(self):
+        with open(self.annotation_file, "r") as f:
+            lines = [l.strip() for l in f]
+        data = []
+        i = 0
+        while i < len(lines):
+            fn = lines[i]; i += 1
+            num = int(lines[i]); i += 1
+            boxes = []
+            for _ in range(num):
+                x,y,w,h = map(int, lines[i].split()[:4]); i += 1
+                if w>0 and h>0: boxes.append((x,y,w,h))
+            if boxes: data.append((fn, boxes))
+        return data
+
+    def encode_target(self, boxes, W, H):
+        t = torch.zeros((self.S,self.S,self.B*5))
+        for x,y,w,h in boxes:
+            cx=(x+w/2)/W; cy=(y+h/2)/H; wn=w/W; hn=h/H
+            gx=int(cx*self.S); gy=int(cy*self.S)
+            if 0<=gx<self.S and 0<=gy<self.S:
+                for b in range(self.B):
+                    if t[gy,gx,b*5+4]==0:
+                        t[gy,gx,b*5:b*5+5]=torch.tensor([cx,cy,wn,hn,1.0])
+                        break
+        return t
+
+    def __getitem__(self, idx):
+        fn, boxes = self.data[idx]
+        img = Image.open(os.path.join(self.images_dir, fn)).convert("RGB")
+        W,H = img.size
+        if self.transform: img = self.transform(img)
+        return img, self.encode_target(boxes,W,H)
+
+    def __len__(self): return len(self.data)

face_detection_checkpoint_epoch40.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fec8368ba1b01937f1f30ed7e01bb641fd9896e70e0b2300daa8103a1658ad84
+size 68801828

inference.py

@@ -0,0 +1,27 @@
+
+import torch, matplotlib.pyplot as plt, matplotlib.patches as patches
+
+def draw_prediction(img, pred, S=20, B=3, thr=0.15):
+    img_np = img.permute(1,2,0).numpy()
+    H,W,_ = img_np.shape
+    fig,ax=plt.subplots(1,figsize=(6,6)); ax.imshow(img_np)
+    c=0
+    for i in range(S):
+        for j in range(S):
+            for b in range(B):
+                off=b*5
+                conf=float(torch.sigmoid(pred[i,j,off+4]))
+                if conf>thr:
+                    cx,cy,w,h = [float(pred[i,j,off+k]) for k in range(4)]
+                    cx*=W; cy*=H; w*=W; h*=H
+                    x=cx-w/2; y=cy-h/2
+                    ax.add_patch(patches.Rectangle((x,y),w,h,edgecolor='red',facecolor='none',linewidth=2))
+                    c+=1
+    ax.set_title(f"{c} boxes"); plt.axis('off'); plt.show()
+
+def load_and_predict(model_path, img_tensor, model):
+    model.load_state_dict(torch.load(model_path, map_location='cpu'))
+    model.eval()
+    with torch.no_grad():
+        out=model(img_tensor.unsqueeze(0)).squeeze(0)
+    draw_prediction(img_tensor, out)

loss.py

@@ -0,0 +1,16 @@
+
+import torch.nn.functional as F
+import torch
+
+def yolo_style_loss(pred, target, S=20, B=3):
+    loss=0
+    for b in range(B):
+        pb = pred[...,b*5:b*5+4]
+        pc = pred[...,b*5+4]
+        tb = target[...,b*5:b*5+4]
+        tc = target[...,b*5+4]
+        pc=torch.sigmoid(pc)
+        mask = tc>0
+        loss += F.smooth_l1_loss(pb[mask], tb[mask], reduction='sum')
+        loss += F.binary_cross_entropy(pc, tc, reduction='sum')
+    return loss

model.py

@@ -0,0 +1,31 @@
+
+import torch
+import torch.nn as nn
+
+class FaceDetectionMultiBoxCNN(nn.Module):
+    def __init__(self, S=20, B=3):
+        super().__init__()
+        self.S = S
+        self.B = B
+        self.output_size = S * S * B * 5
+
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 32, 3, 1, 1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, 1, 1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(128, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(256, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(), nn.MaxPool2d(2),
+        )
+
+        dummy = torch.zeros(1, 3, 320, 320)
+        with torch.no_grad():
+            flatten = self.features(dummy).view(1, -1).shape[1]
+
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(flatten, 1024), nn.ReLU(),
+            nn.Linear(1024, self.output_size)
+        )
+
+    def forward(self, x):
+        return self.classifier(self.features(x)).view(-1, self.S, self.S, self.B * 5)

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+matplotlib
+Pillow