add list of weights

weights/get_weights.py

@@ -1,5 +1,21 @@
+from ast import arg
 import gdown
+import argparse
 
-url = "https://drive.google.com/uc?id=1Bw4gUsRBxy8XZDGchPJ_URQjbHItikjw"
-output = "./resnet_10.pkl"
-gdown.download(url, output, quiet=False)
+parser = argparse.ArgumentParser(description='Download weights model')
+parser.add_argument('--weights', type=str, default='resnet', help='Weights name')
+
+args = parser.parse_args()
+
+def get_weights(weights):
+    weights_list = {
+        'resnet': '1Bw4gUsRBxy8XZDGchPJ_URQjbHItikjw',
+        'resnet18': '1k_v1RrDO6da_NDhBtMZL5c0QSogCmiRn',
+        'vgg11': ''
+    }
+
+    url = f"https://drive.google.com/uc?id={weights_list[weights]}"
+    output = f"./{weights}.pkl"
+    gdown.download(url, output, quiet=False)
+
+get_weights(args.weights)