Update app.py

app.py

@@ -1,157 +1,124 @@
-import gradio as gr
-import torch
 import os
-from torch.utils.data import Dataset, DataLoader
-from torchvision import transforms
-from PIL import Image
-import xml.etree.ElementTree as ET
-import torch.optim as optim
 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
 
-# Device config
+# ----------- 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 -----------
 
-# Custom Dataset
 class FaceMaskDataset(Dataset):
-    def __init__(self, images_dir, annotations_dir, transform=None, resize=(800, 800)):
-        self.images_dir = images_dir
-        self.annotations_dir = annotations_dir
+    def __init__(self, root_dir, transform=None):
+        self.image_paths = []
+        self.labels = []
         self.transform = transform
-        self.resize = resize
-        self.image_files = os.listdir(images_dir)
+        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_files)
+        return len(self.image_paths)
 
     def __getitem__(self, idx):
-        image_path = os.path.join(self.images_dir, self.image_files[idx])
-        image = Image.open(image_path).convert("RGB")
-        image = image.resize(self.resize)
-
-        annotation_path = os.path.join(
-            self.annotations_dir,
-            self.image_files[idx].replace(".jpg", ".xml").replace(".png", ".xml")
-        )
-        if not os.path.exists(annotation_path):
-            print(f"Warning: Annotation file {annotation_path} not found.")
-            return None, None
-
-        boxes, labels = self.load_annotations(annotation_path)
-        if boxes is None or labels is None:
-            return None, None
-
-        target = {'boxes': boxes, 'labels': labels}
+        image = Image.open(self.image_paths[idx]).convert("RGB")
         if self.transform:
             image = self.transform(image)
+        return image, self.labels[idx]
 
-        return image, target
-
-    def load_annotations(self, annotation_path):
-        tree = ET.parse(annotation_path)
-        root = tree.getroot()
-
-        boxes = []
-        labels = []
-        for obj in root.iter('object'):
-            label = obj.find('name').text
-            bndbox = obj.find('bndbox')
-            xmin = float(bndbox.find('xmin').text)
-            ymin = float(bndbox.find('ymin').text)
-            xmax = float(bndbox.find('xmax').text)
-            ymax = float(bndbox.find('ymax').text)
-            boxes.append([xmin, ymin, xmax, ymax])
-            labels.append(1 if label == "mask" else 0)
-
-        if not boxes or not labels:
-            return None, None
-
-        return torch.as_tensor(boxes, dtype=torch.float32), torch.tensor(labels, dtype=torch.int64)
-
-# Placeholder collate function
-def collate_fn(batch):
-    batch = list(filter(lambda x: x[0] is not None, batch))
-    images, targets = zip(*batch)
-    return images, targets
-
-# Dummy get_model function (replace with real model)
-def get_model(num_classes):
-    import torchvision
-    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
-    in_features = model.roi_heads.box_predictor.cls_score.in_features
-    model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)
-    return model
-
-# Validation Function
-def evaluate_model(model, val_loader):
-    model.eval()
-    running_loss = 0.0
-    with torch.no_grad():
-        for images, targets in val_loader:
-            if images is None or targets is None:
-                continue
-            images = [img.to(device) for img in images]
-            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
-            loss_dict = model(images, targets)
-            total_loss = sum(loss for loss in loss_dict.values())
-            running_loss += total_loss.item()
-    return running_loss / len(val_loader)
-
-# Training Function
-def train_model(model, train_loader, val_loader, optimizer, num_epochs=10):
-    for epoch in range(num_epochs):
-        running_loss = 0.0
-        model.train()
-        for images, targets in train_loader:
-            if images is None or targets is None:
-                continue
-            images = [img.to(device) for img in images]
-            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
-            optimizer.zero_grad()
-            loss_dict = model(images, targets)
-            total_loss = sum(loss for loss in loss_dict.values())
-            total_loss.backward()
-            optimizer.step()
-            running_loss += total_loss.item()
-
-        print(f"[Epoch {epoch+1}] Train Loss: {running_loss / len(train_loader):.4f}")
-        val_loss = evaluate_model(model, val_loader)
-        print(f"[Epoch {epoch+1}] Validation Loss: {val_loss:.4f}")
-
-    torch.save(model.state_dict(), "facemask_detector.pth")
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+])
 
-# Main Training Trigger
-def train_from_files_tab():
-    train_zip_path = "train.zip"
-    val_zip_path = "val.zip"
+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)
 
-    if not os.path.exists(train_zip_path) or not os.path.exists(val_zip_path):
-        return "❌ 'train.zip' or 'val.zip' not found in the Files section."
+# ----------- MODEL -----------
 
-    # Extract
-    for zip_path, folder in [(train_zip_path, "train"), (val_zip_path, "val")]:
-        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            zip_ref.extractall(folder)
+model = models.mobilenet_v2(pretrained=True)
+model.classifier[1] = nn.Linear(model.last_channel, 2)
+model = model.to(device)
 
-    transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
-    train_dataset = FaceMaskDataset("train/images", "train/annotations", transform=transform)
-    val_dataset = FaceMaskDataset("val/images", "val/annotations", transform=transform)
+criterion = nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
-    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)
-    val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, collate_fn=collate_fn)
+# ----------- TRAINING -----------
 
-    model = get_model(num_classes=2)
-    model.to(device)
-    optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)
+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"
 
-    train_model(model, train_loader, val_loader, optimizer, num_epochs=5)
-    return "✅ Training complete. Model saved as 'facemask_detector.pth'."
+# ----------- GRADIO APP -----------
 
-# Gradio UI
 iface = gr.Interface(
-    fn=train_from_files_tab,
-    inputs=[],
-    outputs=gr.Textbox(label="Training Output"),
-    title="Face Mask Detector Trainer (from Files Tab)"
+    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()