app.py

import os
import zipfile
from PIL import Image
import torch
import torch.nn as nn
from torchvision import transforms, models
from torch.utils.data import Dataset, DataLoader
import gradio as gr

# ----------- SETUP -----------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# ----------- UNZIP DATA -----------

def unzip_file(zip_path, extract_to):
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to)
        print(f"Extracted {zip_path} to {extract_to}")

unzip_file("train.zip", "./data/train")
unzip_file("val.zip", "./data/val")

# ----------- DATASET -----------

class FaceMaskDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.image_paths = []
        self.labels = []
        self.transform = transform
        for label_name in ['mask', 'no_mask']:
            class_path = os.path.join(root_dir, label_name)
            for img_name in os.listdir(class_path):
                if img_name.endswith(".jpg") or img_name.endswith(".png"):
                    self.image_paths.append(os.path.join(class_path, img_name))
                    self.labels.append(0 if label_name == 'mask' else 1)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, self.labels[idx]

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

train_dataset = FaceMaskDataset("./data/train", transform)
val_dataset = FaceMaskDataset("./data/val", transform)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16)

# ----------- MODEL -----------

model = models.mobilenet_v2(pretrained=True)
model.classifier[1] = nn.Linear(model.last_channel, 2)
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# ----------- TRAINING -----------

def train_model(model, epochs=2):  # keep epochs small for HF Spaces
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        for imgs, labels in train_loader:
            imgs, labels = imgs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(imgs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}")

        # Validation Accuracy
        correct = 0
        total = 0
        model.eval()
        with torch.no_grad():
            for imgs, labels in val_loader:
                imgs, labels = imgs.to(device), labels.to(device)
                outputs = model(imgs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        acc = 100 * correct / total
        print(f"Validation Accuracy: {acc:.2f}%")

train_model(model)
torch.save(model.state_dict(), "face_mask_model.pth")

# ----------- INFERENCE -----------

def predict(image):
    model.eval()
    img = image.convert("RGB")
    img = transform(img).unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = model(img)
        _, predicted = torch.max(outputs.data, 1)
    return "Mask" if predicted.item() == 0 else "No Mask"

# ----------- GRADIO APP -----------

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(source="webcam", tool="editor", type="pil", label="Upload or Webcam"),
    outputs=gr.Label(label="Prediction"),
    live=True,
    title="Face Mask Detection",
    description="Upload or use webcam to detect if a person is wearing a face mask."
)

iface.launch()