xcssgzs commited on
Commit
fd78319
1 Parent(s): c4c6fe1

Upload script.py

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