xcssgzs commited on
Commit
3265605
1 Parent(s): 4239078

Delete script.py

Browse files
Files changed (1) hide show
  1. script.py +0 -90
script.py DELETED
@@ -1,90 +0,0 @@
1
- import os
2
- import pandas as pd
3
- import torch
4
- from PIL import Image
5
- from torchvision import transforms
6
-
7
- from model import efficientnetv2_s as create_model
8
-
9
-
10
- def predict(test_metadata, root_path='/tmp/data/private_testset', output_csv_path='./submission.csv'):
11
-
12
- img_size = {"s": [384, 384], # train_size, val_size
13
- "m": [384, 480],
14
- "l": [384, 480]}
15
- num_model = "s"
16
-
17
- data_transform = transforms.Compose(
18
- [transforms.Resize(img_size[num_model][1]),
19
- transforms.CenterCrop(img_size[num_model][1]),
20
- transforms.ToTensor(),
21
- transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
22
-
23
- id_list = test_metadata['observation_id'].tolist()
24
- img_name_list = test_metadata['filename'].tolist()
25
- print(os.path.abspath(os.path.dirname(__file__)))
26
-
27
- id2classId = dict()
28
- id2prob = dict()
29
- prob_list = list()
30
- classId_list = list()
31
-
32
- for img_name in img_name_list:
33
- img_path = os.path.join(root_path, img_name)
34
- assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
35
- img = Image.open(img_path).convert('RGB')
36
- img = data_transform(img)
37
- img = torch.unsqueeze(img, dim=0)
38
-
39
- with torch.no_grad():
40
- # predict class
41
- output = model(img.to(device)).cpu()
42
- predict = torch.softmax(output, dim=1)
43
- probs, classesId = torch.max(predict, dim=1)
44
- prob = probs.data.numpy().tolist()[0]
45
- classesId = classesId.data.numpy().tolist()[0]
46
- prob_list.append(prob)
47
- classId_list.append(classesId)
48
-
49
- for i, id in enumerate(id_list):
50
- if id not in id2classId.keys():
51
- id2classId[id] = classId_list[i]
52
- id2prob[id] = prob_list[i]
53
- else:
54
- if prob_list[i] > id2prob[id]:
55
- id2classId[id] = classId_list[i]
56
- id2prob[id] = prob_list[i]
57
- classes = list()
58
- for id in id_list:
59
- classes.append(str(id2classId[id]))
60
- test_metadata["class_id"] = classes
61
-
62
- user_pred_df = test_metadata.drop_duplicates("observation_id", keep="first")
63
- user_pred_df[["observation_id", "class_id"]].to_csv(output_csv_path, index=None)
64
-
65
-
66
- if __name__ == '__main__':
67
- import zipfile
68
-
69
- with zipfile.ZipFile("/tmp/data/private_testset.zip", 'r') as zip_ref:
70
- zip_ref.extractall("/tmp/data")
71
- root_path = '/tmp/data/private_testset'
72
-
73
- # root_path = "../../data_set/flower_data/val/n1"
74
-
75
- # json_file = open(json_path, "r")
76
- # index2class = json.load(json_file)
77
-
78
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
79
- # create model
80
- model = create_model(num_classes=1784).to(device)
81
-
82
- # load model weights
83
- model_weight_path = "./efficientNetV2.pth"
84
- model.load_state_dict(torch.load(model_weight_path, map_location=device))
85
- model.eval()
86
-
87
- metadata_file_path = "./SnakeCLEF2024_TestMetadata.csv"
88
- # metadata_file_path = "./test1.csv"
89
- test_metadata = pd.read_csv(metadata_file_path)
90
- predict(test_metadata, root_path)