Upload 4 files

index.py

@@ -0,0 +1,28 @@
+import torch
+from torchvision import transforms
+from PIL import Image, ImageFile
+import pandas as pd
+import os
+import math
+from model import ConvolutionalNet
+from collections import Counter
+from vector_dict import vector_dict
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+model = ConvolutionalNet()
+model.load_state_dict(torch.load('model.pt'))
+
+transform = transforms.Compose([
+  transforms.ToTensor(),
+  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+  transforms.Resize((256, 256))
+])
+
+def get_prediction(path):
+    img = Image.open(path)
+    with torch.no_grad():
+      pred = model(transform(img))
+    return vector_dict[torch.max(pred, 1)[1].item()]
+
+print(get_prediction('data/test/Afghanistan/39841.png'))

model.pt

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

train.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms, datasets, models
+from torchvision.utils import make_grid
+import os
+import time
+from PIL import ImageFile
+import math
+from model import ConvolutionalNet
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+train_transforms = transforms.Compose([
+    transforms.RandomRotation(10),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+test_transform = transforms.Compose([
+  transforms.ToTensor(),
+  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+train_dataset = datasets.ImageFolder(root='./data/train', transform=train_transforms)
+test_dataset = datasets.ImageFolder(root='./data/test', transform=test_transform)
+
+torch.manual_seed(42)
+train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=10, shuffle=False)
+
+class_names = train_dataset.classes
+
+for images, labels in train_loader:
+  break
+
+torch.manual_seed(101)
+model = ConvolutionalNet()
+criterion = nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+start_time = time.time()
+epochs = 5
+
+# BATCH LIMITS
+max_trn_batch = 800
+max_tst_batch = 300
+
+train_losses = []
+test_losses = []
+train_correct = []
+test_correct = []
+
+for epoch in range(epochs):
+  trn_corr = 0
+  tst_corr = 0
+
+  for b, (X_train, y_train) in enumerate(train_loader):
+    # if b == max_trn_batch:
+    #   break
+    
+    y_pred = model(X_train)
+    loss = criterion(y_pred, y_train)
+
+    if b % 200 == 0:
+      print(f"Epoch: {epoch+1}/{epochs}\tBatch: {b+1}\tLoss: {loss.item()}")
+
+    predicted = torch.max(y_pred, 1)[1]
+    batch_corr = (predicted == y_train).sum()
+    trn_corr += batch_corr
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    train_losses.append(loss)
+    train_correct.append(trn_corr)
+  
+  # TEST
+  with torch.no_grad():
+    for b, (X_test, y_test) in enumerate(test_loader):
+      # if b == max_tst_batch:
+      #   break
+      
+      try:
+        y_pred = model(X_test)
+      except:
+        print("Error testing images")
+        continue
+      loss = criterion(y_pred, y_test)
+
+      predicted = torch.max(y_pred, 1)[1]
+      batch_corr = (predicted == y_test).sum()
+      tst_corr += batch_corr
+
+      test_losses.append(loss)
+      test_correct.append(tst_corr)
+
+end_time = time.time()
+total_time = end_time - start_time
+print(f"Time taken: minutes: {math.floor(total_time / 60)} seconds: {math.floor(total_time % 60)}")
+
+torch.save(model.state_dict(), 'model.pt')
+
+plt.plot([x.detach().numpy() for x in train_losses], label='train loss')
+plt.plot(test_losses, label='test loss')
+plt.legend()
+plt.plot()
+
+plt.plot([t/80 for t in train_correct], label='train accuracy')
+plt.plot([t/30 for t in test_correct], label='test accuracy')
+plt.legend()
+plt.plot()
+
+print(f'Accuracy: {100*test_correct[-1].item()/1000}')

vector_dict.py

@@ -0,0 +1 @@
+vector_dict = {0: 'Afghanistan', 1: 'Albania', 2: 'Algeria', 3: 'American Samoa', 4: 'Andorra', 5: 'Angola', 6: 'Anguilla', 7: 'Antarctica', 8: 'Antigua and Barbuda', 9: 'Argentina', 10: 'Armenia', 11: 'Aruba', 12: 'Australia', 13: 'Austria', 14: 'Azerbaijan', 15: 'Bahamas', 16: 'Bahrain', 17: 'Bangladesh', 18: 'Barbados', 19: 'Belarus', 20: 'Belgium', 21: 'Belize', 22: 'Benin', 23: 'Bermuda', 24: 'Bhutan', 25: 'Bolivia', 26: 'Bosnia and Herzegovina', 27: 'Botswana', 28: 'Bouvet Island', 29: 'Brazil', 30: 'British Indian Ocean Territory', 31: 'British Virgin Islands', 32: 'Brunei', 33: 'Bulgaria', 34: 'Burkina Faso', 35: 'Burundi', 36: 'Cambodia', 37: 'Cameroon', 38: 'Canada', 39: 'Cape Verde', 40: 'Caribbean Netherlands', 41: 'Cayman Islands', 42: 'Central African Republic', 43: 'Chad', 44: 'Chile', 45: 'China', 46: 'Christmas Island', 47: 'Cocos (Keeling) Islands', 48: 'Colombia', 49: 'Comoros', 50: 'Cook Islands', 51: 'Costa Rica', 52: 'Croatia', 53: 'Cuba', 54: 'Curaçao', 55: 'Cyprus', 56: 'Czechia', 57: 'DR Congo', 58: 'Denmark', 59: 'Djibouti', 60: 'Dominica', 61: 'Dominican Republic', 62: 'Ecuador', 63: 'Egypt', 64: 'El Salvador', 65: 'Equatorial Guinea', 66: 'Eritrea', 67: 'Estonia', 68: 'Eswatini', 69: 'Ethiopia', 70: 'Falkland Islands', 71: 'Faroe Islands', 72: 'Fiji', 73: 'Finland', 74: 'France', 75: 'French Guiana', 76: 'French Polynesia', 77: 'French Southern and Antarctic Lands', 78: 'Gabon', 79: 'Gambia', 80: 'Georgia', 81: 'Germany', 82: 'Ghana', 83: 'Gibraltar', 84: 'Greece', 85: 'Greenland', 86: 'Grenada', 87: 'Guadeloupe', 88: 'Guam', 89: 'Guatemala', 90: 'Guernsey', 91: 'Guinea', 92: 'Guinea-Bissau', 93: 'Guyana', 94: 'Haiti', 95: 'Heard Island and McDonald Islands', 96: 'Honduras', 97: 'Hong Kong', 98: 'Hungary', 99: 'Iceland', 100: 'India', 101: 'Indonesia', 102: 'Iran', 103: 'Iraq', 104: 'Ireland', 105: 'Isle of Man', 106: 'Israel', 107: 'Italy', 108: 'Ivory Coast', 109: 'Jamaica', 110: 'Japan', 111: 'Jersey', 112: 'Jordan', 113: 'Kazakhstan', 114: 'Kenya', 115: 'Kiribati', 116: 'Kosovo', 117: 'Kuwait', 118: 'Kyrgyzstan', 119: 'Laos', 120: 'Latvia', 121: 'Lebanon', 122: 'Lesotho', 123: 'Liberia', 124: 'Libya', 125: 'Liechtenstein', 126: 'Lithuania', 127: 'Luxembourg', 128: 'Macau', 129: 'Madagascar', 130: 'Malawi', 131: 'Malaysia', 132: 'Maldives', 133: 'Mali', 134: 'Malta', 135: 'Marshall Islands', 136: 'Martinique', 137: 'Mauritania', 138: 'Mauritius', 139: 'Mayotte', 140: 'Mexico', 141: 'Micronesia', 142: 'Moldova', 143: 'Monaco', 144: 'Mongolia', 145: 'Montenegro', 146: 'Montserrat', 147: 'Morocco', 148: 'Mozambique', 149: 'Myanmar', 150: 'Namibia', 151: 'Nauru', 152: 'Nepal', 153: 'Netherlands', 154: 'New Caledonia', 155: 'New Zealand', 156: 'Nicaragua', 157: 'Niger', 158: 'Nigeria', 159: 'Niue', 160: 'Norfolk Island', 161: 'North Korea', 162: 'North Macedonia', 163: 'Northern Mariana Islands', 164: 'Norway', 165: 'Oman', 166: 'Pakistan', 167: 'Palau', 168: 'Palestine', 169: 'Panama', 170: 'Papua New Guinea', 171: 'Paraguay', 172: 'Peru', 173: 'Philippines', 174: 'Pitcairn Islands', 175: 'Poland', 176: 'Portugal', 177: 'Puerto Rico', 178: 'Qatar', 179: 'Republic of the Congo', 180: 'Romania', 181: 'Russia', 182: 'Rwanda', 183: 'Réunion', 184: 'Saint Barthélemy', 185: 'Saint Helena Ascension and Tristan da Cunha', 186: 'Saint Kitts and Nevis', 187: 'Saint Lucia', 188: 'Saint Martin', 189: 'Saint Pierre and Miquelon', 190: 'Saint Vincent and the Grenadines', 191: 'Samoa', 192: 'San Marino', 193: 'Saudi Arabia', 194: 'Senegal', 195: 'Serbia', 196: 'Seychelles', 197: 'Sierra Leone', 198: 'Singapore', 199: 'Sint Maarten', 200: 'Slovakia', 201: 'Slovenia', 202: 'Solomon Islands', 203: 'Somalia', 204: 'South Africa', 205: 'South Georgia', 206: 'South Korea', 207: 'South Sudan', 208: 'Spain', 209: 'Sri Lanka', 210: 'Sudan', 211: 'Suriname', 212: 'Svalbard and Jan Mayen', 213: 'Sweden', 214: 'Switzerland', 215: 'Syria', 216: 'São Tomé and Príncipe', 217: 'Taiwan', 218: 'Tajikistan', 219: 'Tanzania', 220: 'Thailand', 221: 'Timor-Leste', 222: 'Togo', 223: 'Tokelau', 224: 'Tonga', 225: 'Trinidad and Tobago', 226: 'Tunisia', 227: 'Turkey', 228: 'Turkmenistan', 229: 'Turks and Caicos Islands', 230: 'Tuvalu', 231: 'Uganda', 232: 'Ukraine', 233: 'United Arab Emirates', 234: 'United Kingdom', 235: 'United States', 236: 'United States Minor Outlying Islands', 237: 'United States Virgin Islands', 238: 'Uruguay', 239: 'Uzbekistan', 240: 'Vanuatu', 241: 'Vatican City', 242: 'Venezuela', 243: 'Vietnam', 244: 'Wallis and Futuna', 245: 'Western Sahara', 246: 'Yemen', 247: 'Zambia', 248: 'Zimbabwe', 249: 'Åland Islands'}