Model training (#104)

* Update repo structure

* added the whole frontend

* Update repo structure

* Adding object detection model

* added object detection model

* added api

* resolved some conflicts

* model training

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Tanisha Lalwani <[email protected]>

Backend/model_training.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
+from torchvision.models import resnet50
+import cv2
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+
+class ExamDataset(Dataset):
+    def __init__(self, image_paths, annotations, transform=None):
+        self.image_paths = image_paths
+        self.annotations = annotations
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        # Load image
+        image = cv2.imread(self.image_paths[idx])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # Get annotations
+        boxes = self.annotations[idx]['boxes']
+        labels = self.annotations[idx]['labels']
+
+        if self.transform:
+            image = self.transform(image)
+
+        target = {
+            'boxes': torch.FloatTensor(boxes),
+            'labels': torch.LongTensor(labels)
+        }
+
+        return image, target
+
+class InvigilationSystem:
+    def __init__(self):
+        # Initialize FRCNN for student detection and behavior analysis
+        self.frcnn = fasterrcnn_resnet50_fpn(pretrained=True)
+        num_classes = 3  # background, cheating, not_cheating
+        in_features = self.frcnn.roi_heads.box_predictor.cls_score.in_features
+        self.frcnn.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+
+        # Initialize CNN for face recognition
+        self.face_cnn = resnet50(pretrained=True)
+        num_features = self.face_cnn.fc.in_features
+        self.face_cnn.fc = nn.Linear(num_features, len(self.known_faces))
+
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.frcnn.to(self.device)
+        self.face_cnn.to(self.device)
+
+    def train_models(self, train_loader, num_epochs=10):
+        # Training parameters
+        params = [p for p in self.frcnn.parameters() if p.requires_grad]
+        optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
+
+        for epoch in range(num_epochs):
+            self.frcnn.train()
+            total_loss = 0
+
+            for images, targets in train_loader:
+                images = [image.to(self.device) for image in images]
+                targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets]
+
+                loss_dict = self.frcnn(images, targets)
+                losses = sum(loss for loss in loss_dict.values())
+
+                optimizer.zero_grad()
+                losses.backward()
+                optimizer.step()
+
+                total_loss += losses.item()
+
+            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(train_loader):.4f}")
+
+    def process_frame(self, frame):
+        self.frcnn.eval()
+        self.face_cnn.eval()
+
+        # Transform frame
+        transform = torchvision.transforms.Compose([
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        frame_tensor = transform(frame).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            predictions = self.frcnn(frame_tensor)
+
+        # Process predictions
+        boxes = predictions[0]['boxes'].cpu().numpy()
+        scores = predictions[0]['scores'].cpu().numpy()
+        labels = predictions[0]['labels'].cpu().numpy()
+
+        results = []
+        for box, score, label in zip(boxes, scores, labels):
+            if score > 0.5:  # Confidence threshold
+                x1, y1, x2, y2 = box.astype(int)
+                face_crop = frame[y1:y2, x1:x2]
+
+                # Face recognition
+                face_tensor = transform(face_crop).unsqueeze(0).to(self.device)
+                face_prediction = self.face_cnn(face_tensor)
+                student_id = torch.argmax(face_prediction).item()
+
+                results.append({
+                    'box': box,
+                    'score': score,
+                    'is_cheating': label == 1,
+                    'student_id': student_id
+                })
+
+        return results
+
+    def generate_report(self, results):
+        report_data = []
+        for result in results:
+            report_data.append({
+                'timestamp': pd.Timestamp.now(),
+                'student_id': result['student_id'],
+                'confidence': result['score'],
+                'behavior': 'Suspicious' if result['is_cheating'] else 'Normal'
+            })
+
+        df = pd.DataFrame(report_data)
+        df.to_excel('invigilation_report.xlsx', index=False)
+        return df
+
+class FastRCNNPredictor(nn.Module):
+    def __init__(self, in_channels, num_classes):
+        super(FastRCNNPredictor, self).__init__()
+        self.cls_score = nn.Linear(in_channels, num_classes)
+        self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
+
+    def forward(self, x):
+        if x.dim() == 4:
+            torch.flatten(x, start_dim=1)
+        scores = self.cls_score(x)
+        bbox_deltas = self.bbox_pred(x)
+        return scores, bbox_deltas

Backend/proctor_core.py

@@ -1,4 +1,3 @@
-
 import cv2
 import mediapipe as mp
 import numpy as np

Backend/run.py

@@ -1,4 +1,3 @@
-
 import threading as th
 import logging
 import os
@@ -107,3 +106,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+

Backend/train.py

@@ -0,0 +1,58 @@
+import os
+import torch
+from torchvision import transforms
+from model_training import InvigilationSystem, ExamDataset
+
+def prepare_dataset(data_dir):
+    image_paths = []
+    annotations = []
+
+    # Load images and annotations from your data directory
+    for image_file in os.listdir(os.path.join(data_dir, 'images')):
+        if image_file.endswith(('.jpg', '.png')):
+            image_paths.append(os.path.join(data_dir, 'images', image_file))
+
+            # Load corresponding annotation file
+            annotation_file = os.path.join(
+                data_dir, 
+                'annotations', 
+                image_file.replace('.jpg', '.json').replace('.png', '.json')
+            )
+            with open(annotation_file, 'r') as f:
+                annotation = json.load(f)
+                annotations.append(annotation)
+
+    return image_paths, annotations
+
+def main():
+    # Set up data transformations
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+
+    # Prepare dataset
+    data_dir = 'path/to/your/dataset'
+    image_paths, annotations = prepare_dataset(data_dir)
+
+    # Create dataset and dataloader
+    dataset = ExamDataset(image_paths, annotations, transform=transform)
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=2,
+        shuffle=True,
+        collate_fn=lambda x: tuple(zip(*x))
+    )
+
+    # Initialize and train the model
+    system = InvigilationSystem()
+    system.train_models(dataloader, num_epochs=10)
+
+    # Save the trained model
+    torch.save({
+        'frcnn_state_dict': system.frcnn.state_dict(),
+        'face_cnn_state_dict': system.face_cnn.state_dict()
+    }, 'invigilation_model.pth')
+
+if __name__ == '__main__':
+    main()