Update app.py

app.py

@@ -6,10 +6,12 @@ from torchvision import transforms
 from PIL import Image
 import xml.etree.ElementTree as ET
 import torch.optim as optim
-from torch import nn
+import zipfile
 
-# Your model training and evaluation functions (already defined in your previous code)
-# Define the custom dataset
+# Device config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Custom Dataset
 class FaceMaskDataset(Dataset):
     def __init__(self, images_dir, annotations_dir, transform=None, resize=(800, 800)):
         self.images_dir = images_dir
@@ -26,18 +28,19 @@ class FaceMaskDataset(Dataset):
         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"))
-        
+        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} does not exist. Skipping image {self.image_files[idx]}.")
-            return None, None  # Return None if annotation is missing
-        
+            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  # Skip if annotations are invalid
+            return None, None
 
         target = {'boxes': boxes, 'labels': labels}
-
         if self.transform:
             image = self.transform(image)
 
@@ -57,119 +60,98 @@ class FaceMaskDataset(Dataset):
             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)  # "mask" = 1, "no_mask" = 0
+            labels.append(1 if label == "mask" else 0)
 
-        if len(boxes) == 0 or len(labels) == 0:
-            return None, None  # If no boxes or labels, return None
+        if not boxes or not labels:
+            return None, None
 
-        boxes = torch.as_tensor(boxes, dtype=torch.float32)
-        labels = torch.tensor(labels, dtype=torch.int64)
+        return torch.as_tensor(boxes, dtype=torch.float32), torch.tensor(labels, dtype=torch.int64)
 
-        return boxes, labels
+# 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
 
-# Model Training Loop (referred to from previous code)
+# 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):
-        # Training loop
         running_loss = 0.0
         model.train()
         for images, targets in train_loader:
             if images is None or targets is None:
-                continue  # Skip invalid images/annotations
-
-            # Move data to device
-            images = [image.to(device) for image in images]
+                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)
-
-            # Calculate total loss
             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}/{num_epochs}, Loss: {running_loss / len(train_loader)}")
-
-        # Evaluate after every epoch
+        print(f"[Epoch {epoch+1}] Train Loss: {running_loss / len(train_loader):.4f}")
         val_loss = evaluate_model(model, val_loader)
-        print(f"Validation Loss: {val_loss}")
+        print(f"[Epoch {epoch+1}] Validation Loss: {val_loss:.4f}")
 
-# 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  # Skip invalid data
+    torch.save(model.state_dict(), "facemask_detector.pth")
 
-            # Move data to device
-            images = [image.to(device) for image in images]
-            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+# Main Training Trigger
+def train_from_files_tab():
+    train_zip_path = "train.zip"
+    val_zip_path = "val.zip"
 
-            loss_dict = model(images, targets)
+    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."
 
-            # Calculate total loss
-            total_loss = sum(loss for loss in loss_dict.values())
-            running_loss += total_loss.item()
+    # 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)
 
-    return running_loss / len(val_loader)
+    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)
 
-# Function to upload dataset and start training
-def train_on_uploaded_data(train_data, val_data):
-    # Save the uploaded dataset (files)
-    train_data_path = "train_data.zip"
-    val_data_path = "val_data.zip"
-    
-    # Unzip and prepare directories (assuming you upload zip files for simplicity)
-    with open(train_data.name, 'wb') as f:
-        f.write(train_data.read())
-    with open(val_data.name, 'wb') as f:
-        f.write(val_data.read())
-
-    # Extract zip files
-    os.system(f"unzip {train_data_path} -d ./train/")
-    os.system(f"unzip {val_data_path} -d ./val/")
-
-    # Load datasets
-    train_dataset = FaceMaskDataset(
-        images_dir="train/images", 
-        annotations_dir="train/annotations", 
-        transform=transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
-    )
-    val_dataset = FaceMaskDataset(
-        images_dir="val/images", 
-        annotations_dir="val/annotations", 
-        transform=transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
-    )
-    
-    # Dataloaders
-    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
-    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)
-
-    # Train the model
-    model = get_model(num_classes=2)  # Assuming you have a model function
+    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)
     model.to(device)
     optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)
 
-    # Train the model and return feedback
-    train_model(model, train_loader, val_loader, optimizer, num_epochs=10)
-
-    return "Training completed and model saved."
+    train_model(model, train_loader, val_loader, optimizer, num_epochs=5)
+    return "✅ Training complete. Model saved as 'facemask_detector.pth'."
 
-# Create Gradio Interface
+# Gradio UI
 iface = gr.Interface(
-    fn=train_on_uploaded_data,
-    inputs=[
-        gr.File(label="Upload Train Dataset (ZIP)"),
-        gr.File(label="Upload Validation Dataset (ZIP)")
-    ],
-    outputs=gr.Textbox(label="Training Status"),
-    live=True
+    fn=train_from_files_tab,
+    inputs=[],
+    outputs=gr.Textbox(label="Training Output"),
+    title="Face Mask Detector Trainer (from Files Tab)"
 )
 
-# Launch Gradio interface
 iface.launch()