From ccb69e29ea081461ee3a2725dbc1f98dc177d629 Mon Sep 17 00:00:00 2001 From: MANI <122438942+devxMani@users.noreply.github.com> Date: Fri, 8 Nov 2024 20:57:21 +0530 Subject: [PATCH] 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 <145191259+tanishaness@users.noreply.github.com> --- Backend/model_training.py | 145 ++++++++++++++++++++++++++++++++++++++ Backend/proctor_core.py | 1 - Backend/run.py | 2 +- Backend/train.py | 58 +++++++++++++++ 4 files changed, 204 insertions(+), 2 deletions(-) create mode 100644 Backend/model_training.py create mode 100644 Backend/train.py diff --git a/Backend/model_training.py b/Backend/model_training.py new file mode 100644 index 0000000..f97a7ac --- /dev/null +++ b/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 \ No newline at end of file diff --git a/Backend/proctor_core.py b/Backend/proctor_core.py index 76df97c..7e58c24 100644 --- a/Backend/proctor_core.py +++ b/Backend/proctor_core.py @@ -1,4 +1,3 @@ - import cv2 import mediapipe as mp import numpy as np diff --git a/Backend/run.py b/Backend/run.py index 06632e0..a701621 100644 --- a/Backend/run.py +++ b/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() + diff --git a/Backend/train.py b/Backend/train.py new file mode 100644 index 0000000..37cc833 --- /dev/null +++ b/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() \ No newline at end of file