Update app.py

app.py

@@ -1,17 +1,15 @@
-import os
+import gradio as gr
 import torch
+import os
 from torch.utils.data import Dataset, DataLoader
-from PIL import Image
-import torchvision
 from torchvision import transforms
+from PIL import Image
 import xml.etree.ElementTree as ET
 import torch.optim as optim
-import matplotlib.pyplot as plt
-import gradio as gr
-
-# Ensure device is set to GPU if available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+from torch import nn
 
+# Your model training and evaluation functions (already defined in your previous code)
+# Define the custom dataset
 class FaceMaskDataset(Dataset):
     def __init__(self, images_dir, annotations_dir, transform=None, resize=(800, 800)):
         self.images_dir = images_dir
@@ -26,21 +24,17 @@ class FaceMaskDataset(Dataset):
     def __getitem__(self, idx):
         image_path = os.path.join(self.images_dir, self.image_files[idx])
         image = Image.open(image_path).convert("RGB")
-
-        # Resize the image to a fixed size, while maintaining aspect ratio
         image = image.resize(self.resize)
 
-        # Handle both .jpg and .png files
         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 a tuple with None to skip the image/annotation pair
+            return None, None  # Return None if annotation is missing
         
         boxes, labels = self.load_annotations(annotation_path)
-
         if boxes is None or labels is None:
-            return None, None  # Skip this item if annotations are invalid
+            return None, None  # Skip if annotations are invalid
 
         target = {'boxes': boxes, 'labels': labels}
 
@@ -63,66 +57,119 @@ 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)  # Assuming "mask" = 1, "no_mask" = 0
+            labels.append(1 if label == "mask" else 0)  # "mask" = 1, "no_mask" = 0
 
         if len(boxes) == 0 or len(labels) == 0:
-            return None, None  # If no boxes or labels are found, return None
+            return None, None  # If no boxes or labels, return None
 
         boxes = torch.as_tensor(boxes, dtype=torch.float32)
         labels = torch.tensor(labels, dtype=torch.int64)
 
         return boxes, labels
 
+# Model Training Loop (referred to from previous code)
+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
 
-# Define the collate function for DataLoader
-def collate_fn(batch):
-    # Filter out None values and pack the rest into a batch
-    batch = [item for item in batch if item[0] is not None and item[1] is not None]
-    return tuple(zip(*batch))
+            # 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]
 
-# Load your pre-trained model (or initialize if required)
-def load_model():
-    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
-    # Assuming 2 classes: mask and no_mask
-    num_classes = 2  
-    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)
-    model.to(device)
-    return model
+            optimizer.zero_grad()
+            loss_dict = model(images, targets)
 
-# Inference function
-def infer(image):
-    model = load_model()  # Load the model
-    model.eval()
-    
-    # Apply transformations
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),  # Resize all images to 224x224
-        transforms.ToTensor(),
-    ])
+            # Calculate total loss
+            total_loss = sum(loss for loss in loss_dict.values())
+            total_loss.backward()
+            optimizer.step()
+
+            running_loss += total_loss.item()
 
-    image = Image.fromarray(image)
-    image = transform(image).unsqueeze(0).to(device)  # Add batch dimension
+        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(train_loader)}")
 
+        # Evaluate after every epoch
+        val_loss = evaluate_model(model, val_loader)
+        print(f"Validation Loss: {val_loss}")
+
+# Validation function
+def evaluate_model(model, val_loader):
+    model.eval()
+    running_loss = 0.0
     with torch.no_grad():
-        prediction = model(image)
+        for images, targets in val_loader:
+            if images is None or targets is None:
+                continue  # Skip invalid data
 
-    # Get boxes and labels from the predictions
-    boxes = prediction[0]['boxes'].cpu().numpy()
-    labels = prediction[0]['labels'].cpu().numpy()
+            # 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]
 
-    return boxes, labels
+            loss_dict = model(images, targets)
 
-# Gradio interface
-def gradio_interface(image):
-    boxes, labels = infer(image)
+            # Calculate total loss
+            total_loss = sum(loss for loss in loss_dict.values())
+            running_loss += total_loss.item()
+
+    return running_loss / len(val_loader)
+
+# 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"
     
-    # Assuming labels: 0 = no mask, 1 = mask
-    result = {"boxes": boxes, "labels": labels}
-    return result
+    # 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)
 
-# Create Gradio interface
-iface = gr.Interface(fn=gradio_interface, inputs=gr.Image(type="numpy"), outputs="json")
+    # Train the model
+    model = get_model(num_classes=2)  # Assuming you have a model function
+    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."
+
+# Create Gradio Interface
+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
+)
 
 # Launch Gradio interface
 iface.launch()