{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":85769,"databundleVersionId":9709110,"sourceType":"competition"}],"dockerImageVersionId":30775,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118","metadata":{"execution":{"iopub.status.busy":"2024-10-03T19:29:00.8112Z","iopub.execute_input":"2024-10-03T19:29:00.812164Z","iopub.status.idle":"2024-10-03T19:29:12.260274Z","shell.execute_reply.started":"2024-10-03T19:29:00.812104Z","shell.execute_reply":"2024-10-03T19:29:12.259182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nprint(torch.cuda.is_available())","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-03T19:29:12.262369Z","iopub.execute_input":"2024-10-03T19:29:12.262748Z","iopub.status.idle":"2024-10-03T19:29:12.269069Z","shell.execute_reply.started":"2024-10-03T19:29:12.262689Z","shell.execute_reply":"2024-10-03T19:29:12.268054Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install alive_progress matplotlib","metadata":{"execution":{"iopub.status.busy":"2024-10-03T19:29:12.270251Z","iopub.execute_input":"2024-10-03T19:29:12.270552Z","iopub.status.idle":"2024-10-03T19:29:23.809527Z","shell.execute_reply.started":"2024-10-03T19:29:12.270519Z","shell.execute_reply":"2024-10-03T19:29:23.808383Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torchvision import transforms, datasets\nfrom torch.utils.data import DataLoader\nimport os\nfrom alive_progress import alive_bar, config_handler\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport csv\n\n# Configure alive_progress\nconfig_handler.set_global(spinner='pulse', bar='smooth')\n\n# Set up data loaders\ndef create_data_loaders(data_dir, batch_size=32):\n    transform = transforms.Compose([\n        transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ])\n\n    full_dataset = datasets.ImageFolder(root=data_dir, transform=transform)\n    train_size = int(0.8 * len(full_dataset))\n    val_size = len(full_dataset) - train_size\n    train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size])\n\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)\n\n    return train_loader, val_loader\n\n# Custom CNN model\nclass CustomCNN(nn.Module):\n    def __init__(self, num_classes):\n        super(CustomCNN, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(3, 64, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(kernel_size=2, stride=2),\n            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(kernel_size=2, stride=2),\n            nn.Conv2d(128, 256, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(kernel_size=2, stride=2)\n        )\n        self.classifier = nn.Sequential(\n            nn.Linear(256 * 28 * 28, 512),\n            nn.ReLU(inplace=True),\n            nn.Dropout(0.5),\n            nn.Linear(512, num_classes)\n        )\n\n    def forward(self, x):\n        x = self.features(x)\n        x = torch.flatten(x, 1)\n        x = self.classifier(x)\n        return x\n\n# Custom weighted loss function\nclass WeightedCrossEntropyLoss(nn.Module):\n    def __init__(self, weights):\n        super(WeightedCrossEntropyLoss, self).__init__()\n        self.weights = weights\n\n    def forward(self, outputs, targets):\n        loss = nn.functional.cross_entropy(outputs, targets, reduction='none')\n        weighted_loss = loss * self.weights[targets]\n        return weighted_loss.mean()\n\ndef train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device):\n    history = {\n        'train_loss': [], 'val_loss': [],\n        'train_acc': [], 'val_acc': []\n    }\n    \n    for epoch in range(num_epochs):\n        model.train()\n        running_loss = 0.0\n        correct_train = 0\n        total_train = 0\n                \n        with alive_bar(len(train_loader), title=f'Epoch {epoch+1}/{num_epochs}', length=50) as bar:\n            i = 0\n            for inputs, labels in train_loader:\n                inputs, labels = inputs.to(device), labels.to(device)\n                optimizer.zero_grad()\n                outputs = model(inputs)\n                loss = criterion(outputs, labels)\n                loss.backward()\n                optimizer.step()\n                running_loss += loss.item()\n\n                _, predicted = outputs.max(1)\n                total_train += labels.size(0)\n                correct_train += predicted.eq(labels).sum().item()\n                \n                current_train_acc = 100. * correct_train / total_train\n                bar.text(f'Train Acc: {current_train_acc:.2f}%')\n                i += 1\n#                 print(f'{i}/{len(train_loader)}')\n                bar()\n\n        model.eval()\n        val_loss = 0.0\n        correct_val = 0\n        total_val = 0\n\n        with torch.no_grad():\n            for inputs, labels in val_loader:\n                inputs, labels = inputs.to(device), labels.to(device)\n                outputs = model(inputs)\n                loss = criterion(outputs, labels)\n                val_loss += loss.item()\n                _, predicted = outputs.max(1)\n                total_val += labels.size(0)\n                correct_val += predicted.eq(labels).sum().item()\n\n        avg_train_loss = running_loss / len(train_loader)\n        avg_val_loss = val_loss / len(val_loader)\n        train_accuracy = 100. * correct_train / total_train\n        val_accuracy = 100. * correct_val / total_val\n\n        history['train_loss'].append(avg_train_loss)\n        history['val_loss'].append(avg_val_loss)\n        history['train_acc'].append(train_accuracy)\n        history['val_acc'].append(val_accuracy)\n\n        print(f\"Epoch {epoch+1}/{num_epochs}\")\n        print(f\"Train Loss: {avg_train_loss:.4f}, Train Acc: {train_accuracy:.2f}%\")\n        print(f\"Val Loss: {avg_val_loss:.4f}, Val Acc: {val_accuracy:.2f}%\")\n        print(\"-\" * 50)\n\n    return model, history\n\n# Function to plot training history\ndef plot_training_history(history):\n    plt.figure(figsize=(12, 4))\n    \n    plt.subplot(1, 2, 1)\n    plt.plot(history['train_loss'], label='Train Loss')\n    plt.plot(history['val_loss'], label='Validation Loss')\n    plt.title('Model Loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend()\n    \n    plt.subplot(1, 2, 2)\n    plt.plot(history['train_acc'], label='Train Accuracy')\n    plt.plot(history['val_acc'], label='Validation Accuracy')\n    plt.title('Model Accuracy')\n    plt.xlabel('Epoch')\n    plt.ylabel('Accuracy')\n    plt.legend()\n    \n    plt.tight_layout()\n    plt.savefig('training_history.png')\n    plt.close()\n\ndef preprocess_image(image_path):\n    transform = transforms.Compose([\n        transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ])\n    image = Image.open(image_path).convert('RGB')\n    return transform(image).unsqueeze(0)\n\ndef predict_image(model, image_tensor, class_names):\n    with torch.no_grad():\n        outputs = model(image_tensor)\n        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)\n        top5_prob, top5_catid = torch.topk(probabilities, 5)\n        return [(class_names[idx], prob.item()) for idx, prob in zip(top5_catid, top5_prob)]\n\ndef evaluate_model(model, test_dir, class_names):\n    model.eval()\n    image_files = [f for f in os.listdir(test_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]\n\n    csv_path = 'predictions.csv'\n    with open(csv_path, 'w', newline='') as csvfile:\n        writer = csv.writer(csvfile)\n        writer.writerow(['Image Name', 'Class1', 'Class2', 'Class3', 'Class4', 'Class5'])\n\n        with alive_bar(len(image_files), title='Evaluating Images', length=50) as bar:\n            for image_file in image_files:\n                image_path = os.path.join(test_dir, image_file)\n                image_tensor = preprocess_image(image_path)\n                predictions = predict_image(model, image_tensor, class_names)\n                \n                writer.writerow([image_file] + [class_name for class_name, _ in predictions])\n                bar()\n\n    print(f\"Predictions saved to {csv_path}\")\n\ndef main(train_dir, test_dir, num_epochs=10, batch_size=32):\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {device}\")\n\n    train_loader, val_loader = create_data_loaders(train_dir, batch_size)\n    class_names = train_loader.dataset.dataset.classes\n    num_classes = len(class_names)\n    \n    model = CustomCNN(num_classes).to(device)\n\n    class_counts = torch.tensor([len(os.listdir(os.path.join(train_dir, class_name))) for class_name in class_names])\n    class_weights = 1.0 / class_counts\n    class_weights = class_weights / class_weights.sum()\n    class_weights = class_weights.to(device)\n\n    criterion = WeightedCrossEntropyLoss(class_weights)\n    optimizer = optim.Adam(model.parameters(), lr=0.0001)\n\n    print(\"Starting training...\")\n    model, history = train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device)\n\n    print(\"Training completed. Saving model...\")\n    torch.save(model.state_dict(), 'custom_cnn_model.pth')\n    print(\"Model saved successfully.\")\n\n    print(\"Plotting training history...\")\n    plot_training_history(history)\n    print(\"Training history plot saved as 'training_history.png'.\")\n\n    print(\"Evaluating model on test data...\")\n    evaluate_model(model, test_dir, class_names)\n\nif __name__ == \"__main__\":\n    train_directory = \"/kaggle/input/artist-identification/artist_dataset/train\"\n    test_directory = \"/kaggle/input/artist-identification/artist_dataset/test\"\n    main(train_directory, test_directory)","metadata":{"execution":{"iopub.status.busy":"2024-10-03T19:29:23.812412Z","iopub.execute_input":"2024-10-03T19:29:23.812786Z"},"trusted":true},"execution_count":null,"outputs":[]}]}