Merge branch 'IT19972626_StudentConcentration' into 'master'

It19972626 student concentration

See merge request !1

eye.py

+import cv2
+import os
+
+# Load and create faceCascade Classifier
+facePath = os.path.dirname(cv2.__file__) + "/data/haarcascade_frontalface_default.xml"
+faceCascade = cv2.CascadeClassifier(facePath)
+
+eyePath = os.path.dirname(cv2.__file__) + "/data/haarcascade_eye.xml"
+eyeCascade = cv2.CascadeClassifier(eyePath)
+
+out = cv2.VideoWriter('faceEyeDetection.mp4', -1, 20.0, (640, 480))
+video = cv2.VideoCapture(0)
+
+while True:
+    try:
+        ret, frames = video.read()
+        gray = cv2.cvtColor(frames, cv2.COLOR_BGR2GRAY)
+
+        # Face Detection
+        faces = faceCascade.detectMultiScale(gray)
+        for (x, y, w, h) in faces: cv2.rectangle(frames, (x, y), (x + w, y + h), (255, 255, 0), 2)
+
+        # Eye Detection
+        eyes = eyeCascade.detectMultiScale(gray)
+        for (x, y, w, h) in eyes: cv2.rectangle(frames, (x, y), (x + w, y + h), (0, 165, 255), 2)
+
+        # Save and Show the Video
+        out.write(frames)
+        cv2.imshow('WebCam', frames)
+
+        if cv2.waitKey(1) & 0xFF == ord('q'): break
+
+    except Exception as e:
+        print(str(e))
+        break
+
+video.release()
+out.release()
+cv2.destroyAllWindows()
\ No newline at end of file

eyemove_ment/eyetrack.py

+# import face_recognition
+import numpy as np
+import cv2 as cv
+import copy
+from matplotlib import pyplot as plt
+import pyautogui
+import time
+import os
+import cv2
+
+# Load the cascade
+face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
+eye_cascade = cv2.CascadeClassifier("haarcascade_eye.xml")
+
+# Load the video
+cap = cv2.VideoCapture(0)
+
+# Variables to store the last position of the eyes
+eye_pos_x = None
+eye_pos_y = None
+
+while True:
+    # Read the frame
+    _, frame = cap.read()
+
+    # Convert to grayscale
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+    # Detect faces
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
+
+    # Loop over the face detections
+    for (x, y, w, h) in faces:
+        roi_gray = gray[y:y + h, x:x + w]
+        roi_color = frame[y:y + h, x:x + w]
+
+        # Detect eyes
+        eyes = eye_cascade.detectMultiScale(roi_gray)
+
+        # Loop over the eye detections
+        for (ex, ey, ew, eh) in eyes:
+            # Draw a rectangle around the eyes
+            cv2.rectangle(roi_color, (ex, ey), (ex + ew, ey + eh), (0, 255, 0), 2)
+
+            # Check if the eyes are open or closed
+            if ew > 20:
+                # Check if this is the first frame
+                if eye_pos_x is None and eye_pos_y is None:
+                    eye_pos_x = ex + ew / 2
+                    eye_pos_y = ey + eh / 2
+                else:
+                    # Calculate the displacement of the eyes
+                    displacement_x = (ex + ew / 2) - eye_pos_x
+                    displacement_y = (ey + eh / 2) - eye_pos_y
+
+                    # Draw the displacement vector
+                    cv2.arrowedLine(roi_color, (int(eye_pos_x), int(eye_pos_y)),
+                                    (int(eye_pos_x + displacement_x), int(eye_pos_y + displacement_y)), (255, 0, 0), 2)
+
+                    # Check if the eyes are looking to the left or right
+                    if displacement_x < 0:
+                        print("Eyes are looking to the left.")
+                    elif displacement_x > 0:
+                        print("Eyes are looking to the right.")
+
+def maxAndMin(featCoords,mult = 1):
+    adj = 10/mult
+    listX = []
+    listY = []
+    for tup in featCoords:
+        listX.append(tup[0])
+        listY.append(tup[1])
+    maxminList = np.array([min(listX)-adj,min(listY)-adj,max(listX)+adj,max(listY)+adj])
+    print(maxminList)
+    return (maxminList*mult).astype(int), (np.array([sum(listX)/len(listX)-maxminList[0], sum(listY)/len(listY)-maxminList[1]])*mult).astype(int)
+
+def findCircs(img):
+    circles = cv.HoughCircles(img, cv.HOUGH_GRADIENT, 2, 20, param1 = 200, param2 = 50, minRadius=1, maxRadius=40)#, minRadius = 0, maxRadius = 30)
+    # circles = np.uint16(np.around(circles))
+    return circles
+
+def findBlobs(img):
+    params = cv.SimpleBlobDetector_Params()
+    params.minThreshold = 10
+    params.maxThreshold = 200
+
+    # params.filterByColor = True
+    # params.blobColor = 0
+
+    params.filterByArea = True
+    params.maxArea = 3000
+
+    # params.filterByCircularity = True
+    # params.minCircularity = 0.1
+
+    detector = cv.SimpleBlobDetector_create(params)
+    keypoints = detector.detect(img)
+
+    # imkeypoints = cv.drawKeypoints(img, keypoints, np.array([]),
+    #                                (0, 0, 255),
+    #                                cv.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
+    return keypoints
+
+
+
+
+
+def getWebcam(feed=False):
+    webcam = cv.VideoCapture(0)
+    # Frame coordinates go frame[y][x]
+    haventfoundeye = True
+    screenw = 1440
+    screenh = 900
+
+    while True:
+        ret, frame = webcam.read()
+        smallframe = cv.resize(copy.deepcopy(frame), (0,0), fy=.15, fx=.15)
+        smallframe = cv.cvtColor(smallframe, cv.COLOR_BGR2GRAY)
+
+        feats = face_recognition.face_landmarks(smallframe)
+        if len(feats) > 0:
+
+            leBds,leCenter = maxAndMin(feats[0]['left_eye'],mult = 1/.15)
+            # reBds,_ = maxAndMin(feats[0]['right_eye'])
+            # print(leBds)
+
+            left_eye = frame[leBds[1]:leBds[3], leBds[0]:leBds[2]]
+            # right_eye = frame[reBds[1]:reBds[3], reBds[0]:reBds[2]]
+
+            left_eye = cv.cvtColor(left_eye, cv.COLOR_BGR2GRAY)
+            ret, thresh = cv.threshold(left_eye, 50, 255, 0)
+
+            # Find weighted average for center of the eye
+            TMP = 255 - np.copy(thresh)#.astype(int)
+            # TMP = TMP[0:-1, 10:-10]
+            # cv.imshow("tmp", TMP)
+            # TMP = cv.blur(TMP, (3, 3))
+            y = np.sum(TMP, axis=1)
+            x = np.sum(TMP, axis=0)
+            # x = TMP[int(len(TMP)/2)]
+            y = y / len(TMP[0])
+            x = x / len(TMP)
+
+            y = y > np.average(y) + np.std(y)#*1.2
+            x = x > np.average(x) + np.std(x)#*1.2
+
+            try:
+                y = int(np.dot(np.arange(1, len(y) + 1), y) / sum(y))
+            except:
+                y = int(np.dot(np.arange(1, len(y) + 1), y) / 1)
+
+            try:
+                x = int(np.dot(np.arange(1, len(x) + 1), x) / sum(x))
+            except:
+                x = int(np.dot(np.arange(1, len(x) + 1), x) / 1)
+
+            haventfoundeye = False
+
+
+            left_eye = cv.cvtColor(left_eye, cv.COLOR_GRAY2BGR)
+            cv.circle(left_eye, (x, y), 2, (20, 20, 120), 3)
+            cv.circle(left_eye, (int(leCenter[0]), int(leCenter[1])), 2, (120, 20, 20), 3)
+
+
+
+            if feed:
+                cv.imshow('frame', left_eye)
+
+
+
+                if cv.waitKey(1) & 0xFF == ord('q'):
+                    break
+            elif not haventfoundeye:
+                plt.imshow(left_eye)
+                plt.show()
+                return left_eye
+
+def getEye(times = 1,frameShrink = 0.15, coords = (0,0), counterStart = 0, folder = "eyes"):
+    os.makedirs(folder, exist_ok=True)
+    webcam = cv.VideoCapture(0)
+    counter = counterStart
+    ims = []
+
+    while counter < counterStart+times:
+
+        ret, frame = webcam.read()
+        smallframe = cv.resize(copy.deepcopy(frame), (0, 0), fy=frameShrink, fx=frameShrink)
+        smallframe = cv.cvtColor(smallframe, cv.COLOR_BGR2GRAY)
+
+        feats = face_recognition.face_landmarks(smallframe)
+        if len(feats) > 0:
+            leBds, leCenter = maxAndMin(feats[0]['left_eye'], mult=1/frameShrink)
+
+            left_eye = frame[leBds[1]:leBds[3], leBds[0]:leBds[2]]
+            # right_eye = frame[reBds[1]:reBds[3], reBds[0]:reBds[2]]
+
+            left_eye = cv.cvtColor(left_eye, cv.COLOR_BGR2GRAY)
+
+            left_eye = cv.resize(left_eye, dsize=(100, 50))
+
+            # D
+            # isplay the image - DEBUGGING ONLY
+            cv.imshow('frame', left_eye)
+
+            if cv.waitKey(1) & 0xFF == ord('q'):
+                break
+
+            cv.imwrite(
+                folder + "/" + str(coords[0]) + "." + str(coords[1]) + "." + str(
+                    counter) + ".jpg", left_eye)
+            counter += 1
+
+
+# # 1440x900
+# for i in [0,720,1440]:
+#     for j in [0,450,900]:q
+for i in [404,951]:
+    for j in [383,767]:
+        print(i,j)
+        pyautogui.moveTo(i, j)
+        input("Press Enter to continue...")
+        pyautogui.moveTo(i, j)
+        getEye(times = 10, coords=(i,j),counterStart=0, folder = "testeyes")
+

eyemove_ment/face_direction.py

+import cv2
+
+# Load the cascade
+face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
+
+# Load the video
+cap = cv2.VideoCapture(0)
+
+# Get the video dimensions
+width = int(cap.get(3))
+height = int(cap.get(4))
+
+# Define the center of the frame
+center_x = width // 2
+center_y = height // 2
+
+while True:
+    # Read the frame
+    _, frame = cap.read()
+
+    # Convert to grayscale
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+    # Detect faces
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
+
+    # Loop over the face detections
+    for (x, y, w, h) in faces:
+        cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
+        # Get the center of the face
+        face_center_x = x + w // 2
+        face_center_y = y + h // 2
+
+        # Check the direction of the face
+        if face_center_x < center_x - 50:
+            print("Face is looking to the left.")
+            # Send an alert
+            #alert()
+        elif face_center_x > center_x + 50:
+            print("Face is looking to the right.")
+            # Send an alert
+            #alert()
+        else:
+            print("Face is looking straight ahead.")
+
+    # Show the frame
+    cv2.imshow("Face Detection", frame)
+
+    # Exit if the user presses 'q'
+    if cv2.waitKey(1) == ord('q'):
+        break

eyemove_ment/mobile.py

+import cv2
+
+# Load the pre-trained model
+net = cv2.dnn.readNetFromCaffe("mobilenet_iter_73000.caffemodel")
+
+# Load the video
+cap = cv2.VideoCapture(0)
+
+while True:
+    # Read the frame
+    _, frame = cap.read()
+
+    # Get the dimensions of the frame
+    height, width = frame.shape[:2]
+
+    # Create a blob from the frame
+    blob = cv2.dnn.blobFromImage(frame, 0.007843, (width, height), (127.5, 127.5, 127.5), False)
+
+    # Pass the blob through the model
+    net.setInput(blob)
+    detections = net.forward()
+
+    # Loop over the detections
+    for i in range(detections.shape[2]):
+        # Get the confidence of the detection
+        confidence = detections[0, 0, i, 2]
+
+        # Filter out weak detections
+        if confidence > 0.5:
+            # Get the class label of the detection
+            class_id = int(detections[0, 0, i, 1])
+
+            # Check if the class label is a mobile phone
+            if class_id == 7:
+                # Get the bounding box of the detection
+                x1 = int(detections[0, 0, i, 3] * width)
+                y1 = int(detections[0, 0, i, 4] * height)
+                x2 = int(detections[0, 0, i, 5] * width)
+                y2 = int(detections[0, 0, i, 6] * height)
+
+                # Draw a rectangle around the detection
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
+
+    # Show the frame
+    cv2.imshow("Mobile Phone Detection", frame)
+
+    # Exit if the user presses 'q'
+    if cv2.waitKey(1) == ord('q'):
+        break

eyemove_ment/mobiledetectionmodel.py

+import os
+import numpy as np
+import torch
+import glob
+import torch.nn as nn
+from torchvision.transforms import transforms
+from torch.utils.data import DataLoader
+from torch.optim import Adam
+from torch.autograd import Variable
+import torchvision
+import pathlib
+
+# checking for device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Transforms
+transformer = transforms.Compose([
+    transforms.Resize((150, 150)),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),  # 0-255 to 0-1, numpy to tensors
+    transforms.Normalize([0.5, 0.5, 0.5],  # 0-1 to [-1,1] , formula (x-mean)/std
+                         [0.5, 0.5, 0.5])
+])
+
+# Dataloader
+
+# Path for training and testing directory
+fo = open('paths', 'r').read().splitlines()
+train_path = fo[0]
+test_path = fo[1]
+pred_path = fo[2]
+# train_path='/home/r33j4n/Desktop/CNN/Train'
+# test_path='/home/r33j4n/Desktop/CNN/Test'
+
+train_loader = DataLoader(
+    torchvision.datasets.ImageFolder(train_path, transform=transformer),
+    batch_size=64, shuffle=True
+)
+test_loader = DataLoader(
+    torchvision.datasets.ImageFolder(test_path, transform=transformer),
+    batch_size=64, shuffle=True
+)
+
+# categories
+root = pathlib.Path(train_path)
+classes = sorted([j.name.split('/')[-1] for j in root.iterdir()])
+
+
+# CNN Network
+
+
+class ConvNet(nn.Module):
+    def __init__(self, num_classes=8):
+        super(ConvNet, self).__init__()
+
+        # Output size after convolution filter
+        # ((w-f+2P)/s) +1
+
+        # Input shape= (256,3,150,150)
+
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
+        # Shape= (256,12,150,150)
+        self.bn1 = nn.BatchNorm2d(num_features=12)
+        # Shape= (256,12,150,150)
+        self.relu1 = nn.ReLU()
+        # Shape= (256,12,150,150)
+
+        self.pool = nn.MaxPool2d(kernel_size=2)
+        # Reduce the image size be factor 2
+        # Shape= (256,12,75,75)
+
+        self.conv2 = nn.Conv2d(in_channels=12, out_channels=20, kernel_size=3, stride=1, padding=1)
+        # Shape= (256,20,75,75)
+        self.relu2 = nn.ReLU()
+        # Shape= (256,20,75,75)
+
+        self.conv3 = nn.Conv2d(in_channels=20, out_channels=32, kernel_size=3, stride=1, padding=1)
+        # Shape= (256,32,75,75)
+        self.bn3 = nn.BatchNorm2d(num_features=32)
+        # Shape= (256,32,75,75)
+        self.relu3 = nn.ReLU()
+        # Shape= (256,32,75,75)
+
+        self.fc = nn.Linear(in_features=75 * 75 * 32, out_features=num_classes)
+
+        # Feed forwad function
+
+    def forward(self, input):
+        output = self.conv1(input)
+        output = self.bn1(output)
+        output = self.relu1(output)
+
+        output = self.pool(output)
+
+        output = self.conv2(output)
+        output = self.relu2(output)
+
+        output = self.conv3(output)
+        output = self.bn3(output)
+        output = self.relu3(output)
+
+        # Above output will be in matrix form, with shape (256,32,75,75)
+
+        output = output.view(-1, 32 * 75 * 75)
+
+        output = self.fc(output)
+
+        return output
+
+
+model = ConvNet(num_classes=8).to(device)
+# Optmizer and loss function
+optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
+loss_function = nn.CrossEntropyLoss()
+
+num_epochs = 15
+# calculating the size of training and testing images
+train_count = len(glob.glob(train_path + '/**/*.jpeg'))
+test_count = len(glob.glob(test_path + '/**/*.jpeg'))
+
+# Model training and saving best model
+
+best_accuracy = 0.0
+
+for epoch in range(num_epochs):
+
+    # Evaluation and training on training dataset
+    model.train()
+    train_accuracy = 0.0
+    train_loss = 0.0
+
+    for i, (images, labels) in enumerate(train_loader):
+        if torch.cuda.is_available():
+            images = Variable(images.cuda())
+            labels = Variable(labels.cuda())
+
+        optimizer.zero_grad()
+
+        outputs = model(images)
+        loss = loss_function(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        train_loss += loss.cpu().data * images.size(0)
+        _, prediction = torch.max(outputs.data, 1)
+
+        train_accuracy += int(torch.sum(prediction == labels.data))
+
+    train_accuracy = train_accuracy / train_count
+    train_loss = train_loss / train_count
+
+    # Evaluation on testing dataset
+    model.eval()
+
+    test_accuracy = 0.0
+    for i, (images, labels) in enumerate(test_loader):
+        if torch.cuda.is_available():
+            images = Variable(images.cuda())
+            labels = Variable(labels.cuda())
+
+        outputs = model(images)
+        _, prediction = torch.max(outputs.data, 1)
+        test_accuracy += int(torch.sum(prediction == labels.data))
+
+    test_accuracy = test_accuracy / test_count
+
+    print('Epoch: ' + str(epoch) + ' Train Loss: ' + str(train_loss) + ' Train Accuracy: ' + str(
+        train_accuracy) + ' Test Accuracy: ' + str(test_accuracy))
+
+    # Save the best model
+    if test_accuracy > best_accuracy:
+        torch.save(model.state_dict(), 'best_checkpoint.model')
+        best_accuracy = test_accuracy
+
+

mobile.py

+import cv2
+
+# Load the pre-trained model
+net = cv2.dnn.readNetFromCaffe("mobilenet_iter_73000.caffemodel")
+
+# Load the video
+cap = cv2.VideoCapture(0)
+
+while True:
+    # Read the frame
+    _, frame = cap.read()
+
+    # Get the dimensions of the frame
+    height, width = frame.shape[:2]
+
+    # Create a blob from the frame
+    blob = cv2.dnn.blobFromImage(frame, 0.007843, (width, height), (127.5, 127.5, 127.5), False)
+
+    # Pass the blob through the model
+    net.setInput(blob)
+    detections = net.forward()
+
+    # Loop over the detections
+    for i in range(detections.shape[2]):
+        # Get the confidence of the detection
+        confidence = detections[0, 0, i, 2]
+
+        # Filter out weak detections
+        if confidence > 0.5:
+            # Get the class label of the detection
+            class_id = int(detections[0, 0, i, 1])
+
+            # Check if the class label is a mobile phone
+            if class_id == 7:
+                # Get the bounding box of the detection
+                x1 = int(detections[0, 0, i, 3] * width)
+                y1 = int(detections[0, 0, i, 4] * height)
+                x2 = int(detections[0, 0, i, 5] * width)
+                y2 = int(detections[0, 0, i, 6] * height)
+
+                # Draw a rectangle around the detection
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
+
+    # Show the frame
+    cv2.imshow("Mobile Phone Detection", frame)
+
+    # Exit if the user presses 'q'
+    if cv2.waitKey(1) == ord('q'):
+        break

mobiledetectionmodel.py

+import os
+import numpy as np
+import torch
+import glob
+import torch.nn as nn
+from torchvision.transforms import transforms
+from torch.utils.data import DataLoader
+from torch.optim import Adam
+from torch.autograd import Variable
+import torchvision
+import pathlib
+
+# checking for device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Transforms
+transformer = transforms.Compose([
+    transforms.Resize((150, 150)),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),  # 0-255 to 0-1, numpy to tensors
+    transforms.Normalize([0.5, 0.5, 0.5],  # 0-1 to [-1,1] , formula (x-mean)/std
+                         [0.5, 0.5, 0.5])
+])
+
+# Dataloader
+
+# Path for training and testing directory
+fo = open('paths', 'r').read().splitlines()
+train_path = fo[0]
+test_path = fo[1]
+pred_path = fo[2]
+# train_path='/home/r33j4n/Desktop/CNN/Train'
+# test_path='/home/r33j4n/Desktop/CNN/Test'
+
+train_loader = DataLoader(
+    torchvision.datasets.ImageFolder(train_path, transform=transformer),
+    batch_size=64, shuffle=True
+)
+test_loader = DataLoader(
+    torchvision.datasets.ImageFolder(test_path, transform=transformer),
+    batch_size=64, shuffle=True
+)
+
+# categories
+root = pathlib.Path(train_path)
+classes = sorted([j.name.split('/')[-1] for j in root.iterdir()])
+
+
+# CNN Network
+
+
+class ConvNet(nn.Module):
+    def __init__(self, num_classes=8):
+        super(ConvNet, self).__init__()
+
+        # Output size after convolution filter
+        # ((w-f+2P)/s) +1
+
+        # Input shape= (256,3,150,150)
+
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
+        # Shape= (256,12,150,150)
+        self.bn1 = nn.BatchNorm2d(num_features=12)
+        # Shape= (256,12,150,150)
+        self.relu1 = nn.ReLU()
+        # Shape= (256,12,150,150)
+
+        self.pool = nn.MaxPool2d(kernel_size=2)
+        # Reduce the image size be factor 2
+        # Shape= (256,12,75,75)
+
+        self.conv2 = nn.Conv2d(in_channels=12, out_channels=20, kernel_size=3, stride=1, padding=1)
+        # Shape= (256,20,75,75)
+        self.relu2 = nn.ReLU()
+        # Shape= (256,20,75,75)
+
+        self.conv3 = nn.Conv2d(in_channels=20, out_channels=32, kernel_size=3, stride=1, padding=1)
+        # Shape= (256,32,75,75)
+        self.bn3 = nn.BatchNorm2d(num_features=32)
+        # Shape= (256,32,75,75)
+        self.relu3 = nn.ReLU()
+        # Shape= (256,32,75,75)
+
+        self.fc = nn.Linear(in_features=75 * 75 * 32, out_features=num_classes)
+
+        # Feed forwad function
+
+    def forward(self, input):
+        output = self.conv1(input)
+        output = self.bn1(output)
+        output = self.relu1(output)
+
+        output = self.pool(output)
+
+        output = self.conv2(output)
+        output = self.relu2(output)
+
+        output = self.conv3(output)
+        output = self.bn3(output)
+        output = self.relu3(output)
+
+        # Above output will be in matrix form, with shape (256,32,75,75)
+
+        output = output.view(-1, 32 * 75 * 75)
+
+        output = self.fc(output)
+
+        return output
+
+
+model = ConvNet(num_classes=8).to(device)
+# Optmizer and loss function
+optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
+loss_function = nn.CrossEntropyLoss()
+
+num_epochs = 15
+# calculating the size of training and testing images
+train_count = len(glob.glob(train_path + '/**/*.jpeg'))
+test_count = len(glob.glob(test_path + '/**/*.jpeg'))
+
+# Model training and saving best model
+
+best_accuracy = 0.0
+
+for epoch in range(num_epochs):
+
+    # Evaluation and training on training dataset
+    model.train()
+    train_accuracy = 0.0
+    train_loss = 0.0
+
+    for i, (images, labels) in enumerate(train_loader):
+        if torch.cuda.is_available():
+            images = Variable(images.cuda())
+            labels = Variable(labels.cuda())
+
+        optimizer.zero_grad()
+
+        outputs = model(images)
+        loss = loss_function(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        train_loss += loss.cpu().data * images.size(0)
+        _, prediction = torch.max(outputs.data, 1)
+
+        train_accuracy += int(torch.sum(prediction == labels.data))
+
+    train_accuracy = train_accuracy / train_count
+    train_loss = train_loss / train_count
+
+    # Evaluation on testing dataset
+    model.eval()
+
+    test_accuracy = 0.0
+    for i, (images, labels) in enumerate(test_loader):
+        if torch.cuda.is_available():
+            images = Variable(images.cuda())
+            labels = Variable(labels.cuda())
+
+        outputs = model(images)
+        _, prediction = torch.max(outputs.data, 1)
+        test_accuracy += int(torch.sum(prediction == labels.data))
+
+    test_accuracy = test_accuracy / test_count
+
+    print('Epoch: ' + str(epoch) + ' Train Loss: ' + str(train_loss) + ' Train Accuracy: ' + str(
+        train_accuracy) + ' Test Accuracy: ' + str(test_accuracy))
+
+    # Save the best model
+    if test_accuracy > best_accuracy:
+        torch.save(model.state_dict(), 'best_checkpoint.model')
+        best_accuracy = test_accuracy
+
+