Upload 12 files

.gitignore

@@ -0,0 +1,3 @@
+data
+archive.zip
+main.ipynb

app.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+import gradio as gr
+from typing import Tuple, Dict
+from timeit import default_timer as timer
+
+class DogEmotionResNet(nn.Module):
+    def __init__(self, num_classes, weights=None):
+        super().__init__()
+        
+        self.resnet = models.resnet50(weights=weights)
+        for param in self.resnet.parameters():
+            param.requires_grad = False
+        
+        in_features = self.resnet.fc.in_features
+        self.resnet.fc = nn.Linear(in_features, num_classes)
+    
+    def forward(self, x):
+        return self.resnet(x)
+
+def load_model(weights_path: str, num_classes: int) -> DogEmotionResNet:
+    resnet_weights = models.ResNet50_Weights.DEFAULT
+    model = DogEmotionResNet(num_classes=num_classes, weights=resnet_weights)
+    model.load_state_dict(torch.load(weights_path))
+    return model
+
+def load_class_names(file_path: str) -> list:
+    with open(file_path, "r") as f:
+        class_names = [emotion.strip() for emotion in f.readlines()]
+    return class_names
+
+def predict(img) -> Tuple[Dict, float]:
+    """Transforms and performs a prediction on img and returns prediction and time taken."""
+    start_time = timer()
+    
+    img = resnet_transform(img).unsqueeze(0)
+    
+    model.eval()
+    with torch.inference_mode():
+        pred_probs = torch.softmax(model(img), dim=1)
+    
+    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+    
+    pred_time = round(timer() - start_time, 5)
+    return pred_labels_and_probs, pred_time
+
+# Load model and class names
+weights_path = 'dog_emotion_model.pth'
+class_names_file = "classes.txt"
+
+model = load_model(weights_path, num_classes=4)
+class_names = load_class_names(class_names_file)
+resnet_weights = models.ResNet50_Weights.DEFAULT
+resnet_transform = resnet_weights.transforms()
+
+# Gradio Interface
+title = "Dog Emotion Classifier 🐶🎭"
+description = "This app classifies the emotion of a dog in an image into one of four categories: happy, sad, angry, or relaxed."
+article = ""
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Label(num_top_classes=5, label="Predictions"),
+        gr.Number(label="Prediction time (s)"),
+    ],
+    title=title,
+    description=description,
+    article=article,
+)
+
+# Launch the app!
+demo.launch()

classes.txt

@@ -0,0 +1,4 @@
+angry
+happy
+relaxed
+sad

dataset.py

@@ -0,0 +1,27 @@
+import torch
+from torchvision import datasets, transforms, models
+
+def get_datasets(data_dir):
+    resnet_weights = models.ResNet50_Weights.DEFAULT
+    
+    data_transforms = {
+        'train': transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ]),
+        'test': transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+    }
+    
+    image_dataset = datasets.ImageFolder(data_dir, transform=resnet_weights.transforms())
+    
+    train_size = int(0.9 * len(image_dataset))
+    test_size = len(image_dataset) - train_size
+    
+    train_dataset, test_dataset = torch.utils.data.random_split(image_dataset, [train_size, test_size], generator=torch.Generator().manual_seed(42))
+    
+    return train_dataset, test_dataset, image_dataset.classes

dog_emotion.ipynb

@@ -0,0 +1 @@
+{"cells":[{"cell_type":"code","execution_count":2,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.816891Z","iopub.status.busy":"2024-07-13T14:16:59.815968Z","iopub.status.idle":"2024-07-13T14:16:59.821084Z","shell.execute_reply":"2024-07-13T14:16:59.820084Z","shell.execute_reply.started":"2024-07-13T14:16:59.816858Z"},"trusted":true},"outputs":[],"source":["import pandas as pd\n","import numpy as np\n","import matplotlib.pyplot as plt\n","import os"]},{"cell_type":"code","execution_count":3,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.837306Z","iopub.status.busy":"2024-07-13T14:16:59.836599Z","iopub.status.idle":"2024-07-13T14:16:59.844964Z","shell.execute_reply":"2024-07-13T14:16:59.843982Z","shell.execute_reply.started":"2024-07-13T14:16:59.837259Z"},"trusted":true},"outputs":[{"data":{"text/plain":["'2.3.0+cu121'"]},"execution_count":3,"metadata":{},"output_type":"execute_result"}],"source":["import torch\n","from torch import nn, optim\n","from torch.utils.data import DataLoader, Dataset\n","from torchvision import datasets, transforms, models\n","torch.__version__"]},{"cell_type":"code","execution_count":4,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.849583Z","iopub.status.busy":"2024-07-13T14:16:59.849298Z","iopub.status.idle":"2024-07-13T14:16:59.855674Z","shell.execute_reply":"2024-07-13T14:16:59.854751Z","shell.execute_reply.started":"2024-07-13T14:16:59.849561Z"},"trusted":true},"outputs":[{"data":{"text/plain":["'cuda'"]},"execution_count":4,"metadata":{},"output_type":"execute_result"}],"source":["device='cuda' if torch.cuda.is_available() else 'cpu'\n","device"]},{"cell_type":"code","execution_count":9,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.866616Z","iopub.status.busy":"2024-07-13T14:16:59.866153Z","iopub.status.idle":"2024-07-13T14:16:59.889656Z","shell.execute_reply":"2024-07-13T14:16:59.888806Z","shell.execute_reply.started":"2024-07-13T14:16:59.866593Z"},"trusted":true},"outputs":[],"source":["# dir = '/kaggle/input/dog-emotion/Dog Emotion/'\n","dir='data/'"]},{"cell_type":"code","execution_count":6,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.919603Z","iopub.status.busy":"2024-07-13T14:16:59.919302Z","iopub.status.idle":"2024-07-13T14:16:59.928354Z","shell.execute_reply":"2024-07-13T14:16:59.927503Z","shell.execute_reply.started":"2024-07-13T14:16:59.919580Z"},"trusted":true},"outputs":[],"source":["resnet_weights=models.ResNet50_Weights.DEFAULT"]},{"cell_type":"code","execution_count":10,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.930583Z","iopub.status.busy":"2024-07-13T14:16:59.930150Z","iopub.status.idle":"2024-07-13T14:16:59.938564Z","shell.execute_reply":"2024-07-13T14:16:59.937780Z","shell.execute_reply.started":"2024-07-13T14:16:59.930552Z"},"trusted":true},"outputs":[],"source":["train_transforms = transforms.Compose([\n","    transforms.Resize((224, 224)),\n","    transforms.ToTensor(),\n","    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n","])\n","\n","test_transforms = transforms.Compose([\n","    transforms.Resize((224, 224)),\n","    transforms.ToTensor(),\n","    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n","])\n"]},{"cell_type":"code","execution_count":11,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:16:59.941623Z","iopub.status.busy":"2024-07-13T14:16:59.941285Z","iopub.status.idle":"2024-07-13T14:17:01.610438Z","shell.execute_reply":"2024-07-13T14:17:01.609614Z","shell.execute_reply.started":"2024-07-13T14:16:59.941593Z"},"trusted":true},"outputs":[],"source":["# train_dataset = DogEmotionDataset(\n","#     images=X_train.values,\n","#     labels=y_train.values,\n","#     transforms=train_transforms,\n","#     classes=classes\n","# )\n","\n","# test_dataset = DogEmotionDataset(\n","#     images=X_test.values,\n","#     labels=y_test.values,\n","#     transforms=test_transforms,\n","#     classes=classes\n","# )\n","\n","dogdataset = datasets.ImageFolder(dir, transform=resnet_weights.transforms())\n","\n","train_size = int(0.9 * len(dogdataset))\n","test_size = len(dogdataset) - train_size\n","\n","train_dataset, test_dataset = torch.utils.data.random_split(dogdataset, [train_size, test_size],generator=torch.Generator().manual_seed(42))"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:17:01.612559Z","iopub.status.busy":"2024-07-13T14:17:01.612244Z","iopub.status.idle":"2024-07-13T14:17:01.617559Z","shell.execute_reply":"2024-07-13T14:17:01.616669Z","shell.execute_reply.started":"2024-07-13T14:17:01.612534Z"},"trusted":true},"outputs":[],"source":["train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)\n","test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)"]},{"cell_type":"code","execution_count":14,"metadata":{},"outputs":[],"source":["with open('classes.txt', 'w') as f:\n","    for item in dogdataset.classes:\n","        f.write(\"%s\\n\" % item)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:17:01.629372Z","iopub.status.busy":"2024-07-13T14:17:01.629027Z","iopub.status.idle":"2024-07-13T14:17:02.324715Z","shell.execute_reply":"2024-07-13T14:17:02.323774Z","shell.execute_reply.started":"2024-07-13T14:17:01.629341Z"},"trusted":true},"outputs":[],"source":["imgs, labels =next(iter(train_loader))\n","imgs.shape"]},{"cell_type":"code","execution_count":1,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:17:02.327344Z","iopub.status.busy":"2024-07-13T14:17:02.327026Z","iopub.status.idle":"2024-07-13T14:17:02.333456Z","shell.execute_reply":"2024-07-13T14:17:02.332447Z","shell.execute_reply.started":"2024-07-13T14:17:02.327314Z"},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Writing model.py\n"]}],"source":["class DogEmotionResNet(nn.Module):\n","    def __init__(self, num_classes, weights=None):\n","        super().__init__()\n","        \n","        self.resnet = models.resnet50(weights=weights)\n","        for param in self.resnet.parameters():\n","            param.requires_grad = False\n","        \n","        in_features = self.resnet.fc.in_features\n","        self.resnet.fc = nn.Linear(in_features, num_classes)\n","    \n","    def forward(self, x):\n","        return self.resnet(x)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:17:02.334806Z","iopub.status.busy":"2024-07-13T14:17:02.334546Z","iopub.status.idle":"2024-07-13T14:17:02.588742Z","shell.execute_reply":"2024-07-13T14:17:02.587715Z","shell.execute_reply.started":"2024-07-13T14:17:02.334783Z"},"trusted":true},"outputs":[],"source":["torch.manual_seed(42)\n","torch.cuda.manual_seed(42)\n","\n","classes = dogdataset.classes\n","\n","model = DogEmotionResNet(num_classes= len(classes), weights=resnet_weights)\n","model.to(device)\n","\n","# test_output = model(imgs.to(device))"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:17:02.590369Z","iopub.status.busy":"2024-07-13T14:17:02.590030Z","iopub.status.idle":"2024-07-13T14:17:02.602762Z","shell.execute_reply":"2024-07-13T14:17:02.601655Z","shell.execute_reply.started":"2024-07-13T14:17:02.590341Z"},"trusted":true},"outputs":[],"source":["def train(model, train_loader, criterion, optimizer, device):\n","    model.train()\n","    running_loss = 0.0\n","    correct_predictions = 0\n","    total_predictions = 0\n","    \n","    for inputs, labels in train_loader:\n","        inputs, labels = inputs.to(device), labels.to(device)\n","        \n","        optimizer.zero_grad()\n","        \n","        outputs = model(inputs)\n","        loss = criterion(outputs, labels)\n","        loss.backward()\n","        optimizer.step()\n","        \n","        running_loss += loss.item()\n","        predicted_labels = outputs.argmax(dim=1)\n","        correct_predictions += (predicted_labels == labels).sum().item()\n","        total_predictions += labels.size(0)\n","        \n","    train_loss = running_loss / len(train_loader)\n","    train_accuracy = correct_predictions / total_predictions\n","    \n","    return train_loss, train_accuracy\n","\n","def validate(model, test_loader, criterion, device):\n","    model.eval()\n","    running_loss = 0.0\n","    correct_predictions = 0\n","    total_predictions = 0\n","    \n","    with torch.inference_mode():\n","        for inputs, labels in test_loader:\n","            inputs, labels = inputs.to(device), labels.to(device)\n","            \n","            outputs = model(inputs)\n","            loss = criterion(outputs, labels)\n","            \n","            running_loss += loss.item()\n","            predicted_labels = outputs.argmax(dim=1)\n","            correct_predictions += (predicted_labels == labels).sum().item()\n","            total_predictions += labels.size(0)\n","            \n","    test_loss = running_loss / len(test_loader)\n","    test_accuracy = correct_predictions / total_predictions\n","    \n","    return test_loss, test_accuracy\n","        "]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:17:02.604726Z","iopub.status.busy":"2024-07-13T14:17:02.604323Z","iopub.status.idle":"2024-07-13T14:20:47.515463Z","shell.execute_reply":"2024-07-13T14:20:47.514558Z","shell.execute_reply.started":"2024-07-13T14:17:02.604682Z"},"trusted":true},"outputs":[],"source":["from tqdm.auto import tqdm\n","\n","EPOCHS = 30\n","criterion = nn.CrossEntropyLoss()\n","optimizer = optim.Adam(model.parameters(), lr=0.002)\n","\n","train_losses = []\n","train_accuracies = []\n","test_losses = []\n","test_accuracies = []\n","\n","for epoch in tqdm(range(EPOCHS)):\n","    train_loss, train_accuracy = train(model, train_loader, criterion, optimizer, device)\n","    test_loss, test_accuracy = validate(model, test_loader, criterion, device)\n","    \n","    train_losses.append(train_loss)\n","    train_accuracies.append(train_accuracy)\n","    test_losses.append(test_loss)\n","    test_accuracies.append(test_accuracy)\n","    \n","    print(f'Epoch: {epoch+1}/{EPOCHS}, Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}, Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}')\n","    "]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:22:05.329623Z","iopub.status.busy":"2024-07-13T14:22:05.328737Z","iopub.status.idle":"2024-07-13T14:22:05.554562Z","shell.execute_reply":"2024-07-13T14:22:05.553621Z","shell.execute_reply.started":"2024-07-13T14:22:05.329587Z"},"trusted":true},"outputs":[],"source":["plt.figure(figsize=(12, 6))\n","plt.subplot(1, 2, 1)\n","plt.plot(train_losses, label='train loss')\n","plt.plot(test_losses, label='test loss')\n","plt.xlabel('Epochs')\n","plt.ylabel('Loss')\n","plt.legend()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:22:05.557018Z","iopub.status.busy":"2024-07-13T14:22:05.556561Z","iopub.status.idle":"2024-07-13T14:22:05.748228Z","shell.execute_reply":"2024-07-13T14:22:05.747325Z","shell.execute_reply.started":"2024-07-13T14:22:05.556984Z"},"trusted":true},"outputs":[],"source":["plt.subplot(1, 2, 2)\n","plt.plot(train_accuracies, label='train accuracy')\n","plt.plot(test_accuracies, label='test accuracy')\n","plt.xlabel('Epochs')\n","plt.ylabel('Accuracy')\n","plt.legend()\n","plt.show()"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:20:47.969342Z","iopub.status.busy":"2024-07-13T14:20:47.968739Z","iopub.status.idle":"2024-07-13T14:20:48.069413Z","shell.execute_reply":"2024-07-13T14:20:48.068345Z","shell.execute_reply.started":"2024-07-13T14:20:47.969309Z"},"trusted":true},"outputs":[],"source":["torch.save(model.state_dict(), 'dog_emotion_model.pth')"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:20:48.071053Z","iopub.status.busy":"2024-07-13T14:20:48.070747Z","iopub.status.idle":"2024-07-13T14:20:48.343977Z","shell.execute_reply":"2024-07-13T14:20:48.343058Z","shell.execute_reply.started":"2024-07-13T14:20:48.071027Z"},"trusted":true},"outputs":[],"source":["model = DogEmotionResNet(num_classes=len(classes))\n","model.load_state_dict(torch.load('dog_emotion_model.pth'))\n","model.to(device)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:22:26.567754Z","iopub.status.busy":"2024-07-13T14:22:26.566902Z","iopub.status.idle":"2024-07-13T14:22:30.940835Z","shell.execute_reply":"2024-07-13T14:22:30.940029Z","shell.execute_reply.started":"2024-07-13T14:22:26.567714Z"},"trusted":true},"outputs":[],"source":["def predict(model, test_loader, device):\n","    model.eval()\n","    predictions = []\n","    \n","    with torch.inference_mode():\n","        for inputs, labels in test_loader:\n","            inputs = inputs.to(device)\n","            outputs = model(inputs)\n","            predicted_labels = outputs.argmax(dim=1)\n","            \n","            predictions.extend(predicted_labels.cpu().numpy())\n","            \n","    return predictions\n","\n","predictions = predict(model, test_loader, device)"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:22:30.942858Z","iopub.status.busy":"2024-07-13T14:22:30.942507Z","iopub.status.idle":"2024-07-13T14:22:30.949446Z","shell.execute_reply":"2024-07-13T14:22:30.948521Z","shell.execute_reply.started":"2024-07-13T14:22:30.942825Z"},"trusted":true},"outputs":[],"source":["predictions[:10]"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2024-07-13T14:22:33.724471Z","iopub.status.busy":"2024-07-13T14:22:33.723796Z","iopub.status.idle":"2024-07-13T14:22:37.403463Z","shell.execute_reply":"2024-07-13T14:22:37.402489Z","shell.execute_reply.started":"2024-07-13T14:22:33.724441Z"},"trusted":true},"outputs":[],"source":["from sklearn.metrics import classification_report\n","\n","print(classification_report([i[1] for i in test_dataset], predictions))\n"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":[]}],"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"datasetId":2882322,"sourceId":4969612,"sourceType":"datasetVersion"}],"dockerImageVersionId":30747,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.9.13"}},"nbformat":4,"nbformat_minor":4}

dog_emotion_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80a013e6847088db2d44c3dfbc9ac5391e5932f612e30a8b330662c302d8e9ea
+size 94389082

model.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+from torchvision import models
+
+class DogEmotionResNet(nn.Module):
+    def __init__(self, num_classes, weights=None):
+        super().__init__()
+        
+        self.resnet = models.resnet50(weights=weights)
+        for param in self.resnet.parameters():
+            param.requires_grad = False
+        
+        in_features = self.resnet.fc.in_features
+        self.resnet.fc = nn.Linear(in_features, num_classes)
+    
+    def forward(self, x):
+        return self.resnet(x)

readme.md

@@ -0,0 +1,28 @@
+# Dog Emotion Classifier 🐶🎭
+
+This project uses PyTorch for model development and Gradio for the user interface to classify the emotions of dogs from images. The classifier identifies emotions as happy, sad, angry, or relaxed.
+
+## Features
+
+- **Model Architecture:** Utilizes a pre-trained ResNet50 model.
+- **User Interface:** Interactive web interface powered by Gradio.
+- **Training and Evaluation:** Scripts for training and evaluating the model on custom datasets.
+
+## Dataset
+
+- The model was trained on the https://www.kaggle.com/datasets/danielshanbalico/dog-emotion dataset.
+
+## Installation
+
+### Prerequisites
+
+- Python 3.7 or higher
+- PyTorch
+- Gradio
+- Other dependencies listed in `requirements.txt`
+
+### Clone the Repository
+
+```bash
+git clone https://github.com/Yuval728/dog-emotion-classifier.git
+cd dog-emotion-classifier

requirements.txt

@@ -0,0 +1,9 @@
+numpy
+torch
+torchvision
+Pillow
+gradio
+pandas
+scikit-learn
+matplotlib
+tqdm

test.py

@@ -0,0 +1,22 @@
+import torch
+from torch.utils.data import DataLoader
+from sklearn.metrics import classification_report
+
+from dataset import get_datasets
+from model import DogEmotionResNet
+from utils import predict
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+data_dir = 'data/'
+
+_, test_dataset, classes = get_datasets(data_dir)
+
+test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
+
+model = DogEmotionResNet(num_classes=len(classes))
+model.load_state_dict(torch.load('dog_emotion_model.pth'))
+model.to(device)
+
+predictions = predict(model, test_loader, device)
+
+print(classification_report([label for _, label in test_dataset], predictions))

train.py

@@ -0,0 +1,66 @@
+import torch
+from torch import nn, optim
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+import matplotlib.pyplot as plt
+from torchvision import models
+from dataset import get_datasets
+from model import DogEmotionResNet
+from utils import train, validate
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+data_dir = 'data/'
+
+train_dataset, test_dataset, classes = get_datasets(data_dir)
+
+train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
+
+with open('classes.txt', 'w') as f:
+    for item in classes:
+        f.write("%s\n" % item)
+
+torch.manual_seed(42)
+torch.cuda.manual_seed(42)
+
+resnet_weights = models.ResNet50_Weights.DEFAULT
+model = DogEmotionResNet(num_classes=len(classes), weights=resnet_weights)
+model.to(device)
+
+EPOCHS = 30
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.002)
+
+train_losses = []
+train_accuracies = []
+test_losses = []
+test_accuracies = []
+
+for epoch in tqdm(range(EPOCHS)):
+    train_loss, train_accuracy = train(model, train_loader, criterion, optimizer, device)
+    test_loss, test_accuracy = validate(model, test_loader, criterion, device)
+    
+    train_losses.append(train_loss)
+    train_accuracies.append(train_accuracy)
+    test_losses.append(test_loss)
+    test_accuracies.append(test_accuracy)
+    
+    print(f'Epoch: {epoch+1}/{EPOCHS}, Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}, Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}')
+
+plt.figure(figsize=(12, 6))
+plt.subplot(1, 2, 1)
+plt.plot(train_losses, label='train loss')
+plt.plot(test_losses, label='test loss')
+plt.xlabel('Epochs')
+plt.ylabel('Loss')
+plt.legend()
+
+plt.subplot(1, 2, 2)
+plt.plot(train_accuracies, label='train accuracy')
+plt.plot(test_accuracies, label='test accuracy')
+plt.xlabel('Epochs')
+plt.ylabel('Accuracy')
+plt.legend()
+plt.show()
+
+torch.save(model.state_dict(), 'dog_emotion_model.pth')

utils.py

@@ -0,0 +1,63 @@
+import torch
+
+def train(model, train_loader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    correct_predictions = 0
+    total_predictions = 0
+    
+    for inputs, labels in train_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+        
+        optimizer.zero_grad()
+        
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        
+        running_loss += loss.item()
+        predicted_labels = outputs.argmax(dim=1)
+        correct_predictions += (predicted_labels == labels).sum().item()
+        total_predictions += labels.size(0)
+        
+    train_loss = running_loss / len(train_loader)
+    train_accuracy = correct_predictions / total_predictions
+    
+    return train_loss, train_accuracy
+
+def validate(model, test_loader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    correct_predictions = 0
+    total_predictions = 0
+    
+    with torch.inference_mode():
+        for inputs, labels in test_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            
+            running_loss += loss.item()
+            predicted_labels = outputs.argmax(dim=1)
+            correct_predictions += (predicted_labels == labels).sum().item()
+            total_predictions += labels.size(0)
+            
+    test_loss = running_loss / len(test_loader)
+    test_accuracy = correct_predictions / total_predictions
+    
+    return test_loss, test_accuracy
+
+def predict(model, test_loader, device):
+    model.eval()
+    predictions = []
+    
+    with torch.inference_mode():
+        for inputs, _ in test_loader:
+            inputs = inputs.to(device)
+            outputs = model(inputs)
+            predicted_labels = outputs.argmax(dim=1)
+            predictions.extend(predicted_labels.cpu().numpy())
+            
+    return predictions