Update app.py

app.py

@@ -1,124 +1,140 @@
+import gradio as gr
 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
+from torchvision import transforms
+from PIL import Image
+import xml.etree.ElementTree as ET
+import torchvision.models.detection
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
-# ----------- 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 -----------
+print(f"Using device: {device}")
 
+# Dataset class
 class FaceMaskDataset(Dataset):
-    def __init__(self, root_dir, transform=None):
-        self.image_paths = []
-        self.labels = []
+    def __init__(self, images_dir, annotations_dir, transform=None, resize=(800, 800)):
+        self.images_dir = images_dir
+        self.annotations_dir = annotations_dir
         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)
+        self.resize = resize
+        self.image_files = [f for f in os.listdir(images_dir) if f.endswith(('.jpg', '.png'))]
 
     def __len__(self):
-        return len(self.image_paths)
+        return len(self.image_files)
 
     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)
+        image_path = os.path.join(self.images_dir, self.image_files[idx])
+        image = Image.open(image_path).convert("RGB")
+        image = image.resize(self.resize)
 
-# ----------- MODEL -----------
+        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):
+            return None, None
 
-model = models.mobilenet_v2(pretrained=True)
-model.classifier[1] = nn.Linear(model.last_channel, 2)
-model = model.to(device)
+        boxes, labels = self.load_annotations(annotation_path)
+        if boxes is None or labels is None:
+            return None, None
 
-criterion = nn.CrossEntropyLoss()
-optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+        target = {'boxes': boxes, 'labels': labels}
 
-# ----------- TRAINING -----------
+        if self.transform:
+            image = self.transform(image)
 
-def train_model(model, epochs=2):  # keep epochs small for HF Spaces
-    for epoch in range(epochs):
+        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
+
+        boxes = torch.tensor(boxes, dtype=torch.float32)
+        labels = torch.tensor(labels, dtype=torch.int64)
+
+        return boxes, labels
+
+def collate_fn(batch):
+    batch = [b for b in batch if b[0] is not None and b[1] is not None]
+    images, targets = zip(*batch)
+    return list(images), list(targets)
+
+def get_model(num_classes):
+    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 = FastRCNNPredictor(in_features, num_classes)
+    return model
+
+def extract_zip(zip_file, extract_to):
+    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+        zip_ref.extractall(extract_to)
+
+def train_model(train_zip, val_zip):
+    extract_zip(train_zip, './data/train')
+    extract_zip(val_zip, './data/val')
+
+    transform = transforms.Compose([transforms.ToTensor()])
+
+    train_dataset = FaceMaskDataset(
+        images_dir='./data/train/train/images',
+        annotations_dir='./data/train/train/annotations',
+        transform=transform
+    )
+    val_dataset = FaceMaskDataset(
+        images_dir='./data/val/val/images',
+        annotations_dir='./data/val/val/annotations',
+        transform=transform
+    )
+
+    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)
+
+    model = get_model(num_classes=2).to(device)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)
+
+    for epoch in range(3):  # Reduce for demo
         model.train()
-        running_loss = 0.0
-        for imgs, labels in train_loader:
-            imgs, labels = imgs.to(device), labels.to(device)
+        total_loss = 0
+        for images, targets in train_loader:
+            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()
-            outputs = model(imgs)
-            loss = criterion(outputs, labels)
+            loss_dict = model(images, targets)
+            loss = sum(loss for loss in loss_dict.values())
             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 -----------
 
+            total_loss += loss.item()
+
+        print(f"Epoch {epoch+1}, Loss: {total_loss / len(train_loader)}")
+
+    torch.save(model.state_dict(), "model.pth")
+    return "Training completed. Model saved as model.pth"
+
+# Gradio upload interface
 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."
+    fn=train_model,
+    inputs=[
+        gr.File(label="Upload Train ZIP"),
+        gr.File(label="Upload Val ZIP")
+    ],
+    outputs="text"
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()