import torch
from torchvision import models
import requests
import json


def get_pretrained_models():
    # 下载预训练模型并和保存在本地
    torch.save(models.alexnet(weights='DEFAULT'), './models/PublicPretrained/alexnet.pth')
    torch.save(models.vgg16(weights='DEFAULT'), './models/PublicPretrained/vgg16.pth')
    torch.save(models.googlenet(weights='DEFAULT'), './models/PublicPretrained/googlenet.pth')
    torch.save(models.resnet18(weights='DEFAULT'), './models/PublicPretrained/resnet18.pth')
    torch.save(models.resnet50(weights='DEFAULT'), './models/PublicPretrained/resnet50.pth')
    torch.save(models.inception_v3(weights='DEFAULT'), './models/PublicPretrained/inception_v3.pth')
    torch.save(models.densenet161(weights='DEFAULT'), './models/PublicPretrained/densenet161.pth')
    torch.save(models.resnext50_32x4d(weights='DEFAULT'), './models/PublicPretrained/resnext50_32x4d.pth')
    torch.save(models.mobilenet_v2(weights='DEFAULT'), './models/PublicPretrained/mobilenet_v2.pth')


def get_imagenet_labels():
    # 下载 ImageNet 类别到本地
    url = 'https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json'
    response = requests.get(url)
    labels = response.json()

    with open('./configs/imagenet_labels.json', 'w') as f:
        json.dump(labels, f)