Upload script.py

script.py

@@ -8,7 +8,7 @@ from torchvision import transforms
 from model import resnet101
 
 
-def predict(test_metadata, index2class, root_path='/tmp/data/private_testset', output_csv_path='./submission.csv'):
+def predict(test_metadata, root_path='/tmp/data/private_testset', output_csv_path='./submission.csv'):
 
     data_transform = transforms.Compose(
         [transforms.Resize(256),
@@ -17,7 +17,8 @@ def predict(test_metadata, index2class, root_path='/tmp/data/private_testset', o
          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
 
     # load image
-    img_name_list = ["../tulip.jpg", "../rose.jpg"]
+    # img_name_list = ["1163.jpg", "1164.jpg"]
+    # id_list = [1, 2]
     id_list = test_metadata['observation_id'].tolist()
     img_name_list = test_metadata['filename'].tolist()
     img_list = []
@@ -51,7 +52,7 @@ def predict(test_metadata, index2class, root_path='/tmp/data/private_testset', o
                     id2prob[id] = probs[i]
         classes = list()
         for id in id_list:
-            classes.append(index2class[str(id2classId[id])])
+            classes.append(str(id2classId[id]))
         test_metadata["class_id"] = classes
 
         user_pred_df = test_metadata.drop_duplicates("observation_id", keep="first")
@@ -72,8 +73,8 @@ if __name__ == '__main__':
     json_path = './class_indices.json'
     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
 
-    json_file = open(json_path, "r")
-    index2class = json.load(json_file)
+    # json_file = open(json_path, "r")
+    # index2class = json.load(json_file)
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     # create model
@@ -90,4 +91,4 @@ if __name__ == '__main__':
     metadata_file_path = "./SnakeCLEF2024_TestMetadata.csv"
     # metadata_file_path = "./test.csv"
     test_metadata = pd.read_csv(metadata_file_path)
-    predict(test_metadata, index2class, root_path)
+    predict(test_metadata, root_path)