Upload 3 files

mush_classifier_20230801.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:851cac064760eab030275140b81d7e04595a6bb283154f87f72c1b6bc5cf6166
+size 94141978

mushroom_class_load_predict.py

@@ -0,0 +1,73 @@
+from torchvision import transforms
+from PIL import Image
+from PIL import ImageFile
+from torch.utils.data import Dataset,DataLoader
+from transformers import AutoImageProcessor, BitModel, AdamW
+import torch
+from datasets import load_dataset
+from torch import Tensor, nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+import os
+import numpy as np
+from sklearn import metrics, model_selection
+from collections import Counter
+
+# model class
+class mush_root_model(torch.nn.Module):
+    def __init__(self, num_labels=1):
+        super(mush_root_model, self).__init__()
+        self.model = BitModel.from_pretrained("google/bit-50")
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(2048, num_labels),
+        )
+    
+    def forward(self, input):
+        outputs = self.model(**input).pooler_output
+        #print(outputs.shape)
+        logits = self.classifier(outputs)
+        return logits
+
+# load model
+model_path="/kaggle/input/mush-room-model-class/mush_classifier_20230801.pth"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = mush_root_model(num_labels=9)
+model.load_state_dict(torch.load(model_path))
+model.to(device)
+image_processor = AutoImageProcessor.from_pretrained("google/bit-50")
+
+# label setting
+labels = ['Amanita', 'Suillus', 'Boletus', 'Lactarius', 'Agaricus', 'Hygrocybe', 'Cortinarius', 'Russula', 'Entoloma']
+toxic_labels = {'Amanita': 1, 'Suillus': 0, 'Boletus': 0, 'Lactarius': 0, 'Agaricus': 0, 'Hygrocybe': 1, 'Cortinarius': 0, 'Russula': 0, 'Entoloma': 1}
+mushroom_address_list = [
+              "Amanita毒蝇伞,伞菌目,鹅膏菌科,鹅膏菌属,主要分布于我国黑龙江、吉林、四川、西藏、云南等地,有毒",
+              "Suillus乳牛肝菌,牛肝菌目,乳牛肝菌科,乳牛肝菌属,分布于吉林、辽宁、山西、安徽、江西、浙江、湖南、四川、贵州等地,无毒",
+              "Boletus丽柄牛肝菌,伞菌目,牛肝菌科,牛肝菌属,分布于云南、陕西、甘肃、西藏等地,有毒",
+              "Lactarius松乳菇,红菇目,红菇科,乳菇属,广泛分布于亚热带松林地,无毒",
+              "Agaricus双孢蘑菇,伞菌目,蘑菇科,蘑菇属,广泛分布于北半球温带,无毒",
+              "Hygrocybe浅黄褐湿伞,伞菌目,蜡伞科,湿伞属,分布于香港(见于松仔园),有毒",
+              "Cortinarius掷丝膜菌,伞菌目,丝膜菌科,丝膜菌属,分布于湖南等地(夏秋季在山毛等阔叶林地上生长)",
+              "Russula褪色红菇,伞菌目,红菇科,红菇属,分布于河北、吉林、四川、江苏、西藏等地,无毒",
+              "Entoloma霍氏粉褶菌,伞菌目,粉褶菌科,粉褶菌属,主要分布于新西兰北岛和南岛西部,有毒",
+              ]
+def image_process(image_path):
+    image = Image.open(image_path)
+    image_pt = image_processor(image,return_tensors="pt")
+    return image_pt
+
+def predict(image_path):
+    image_pt = image_process(image_path)
+    images = image_pt.to(device)
+    #print(images['pixel_values'].shape)
+    outputs = torch.squeeze(model(images))
+    output = torch.sigmoid(outputs).cpu().detach().numpy().tolist()
+    label_id = np.argmax(output, axis=-1)
+    label_score = output[label_id]
+    return label_id, label_score
+
+
+image_path="/kaggle/input/mush-dataset/dataset1/无毒类/Cortinarius/000_Pw3qUBVmwN8.jpg"
+mushroom_class, confidence = predict(image_path)
+toxic = toxic_labels[labels[mushroom_class]]
+address = mushroom_address_list[mushroom_class]
+print(f"the class of mushroom is { labels[mushroom_class]}, its confidence is {confidence} and it is {bool(toxic)} toxic") 

mushroom_class_train.py

@@ -0,0 +1,155 @@
+from torchvision import transforms
+from PIL import Image
+from PIL import ImageFile
+from torch.utils.data import Dataset,DataLoader
+from transformers import AutoImageProcessor, BitModel, AdamW
+import torch
+from datasets import load_dataset
+from torch import Tensor, nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+import os
+import numpy as np
+from sklearn import metrics, model_selection
+from collections import Counter
+
+
+# cofig_size
+configs = dict()
+configs['batch_size'] = 32
+configs['EPOCHS'] = 10
+configs['LEARNING_RATE'] = 2e-5
+configs['split_size'] = 0.1
+
+# get class label
+label_info = set()
+dataset_path = "/kaggle/input/mush-dataset/dataset1"
+def recur_label(dataset_path):        
+    if os.path.isfile(dataset_path):
+        if ".docx" in dataset_path:
+            return 
+        label_info.add(dataset_path.split("/")[-2])
+    else:
+        for file_name in os.listdir(dataset_path):
+            new_path = dataset_path + "/" + file_name
+            recur_label(new_path)
+    return 
+recur_label(dataset_path)
+print(label_info)
+dict_label = {}
+for i, val in enumerate(label_info):
+    dict_label[val] = i
+print(dict_label)
+
+# deal with dataset, each image_path and label will be in the "all_data" and "all_label" 
+all_data = []
+all_label = []
+dataset_path = "/kaggle/input/mush-dataset/dataset1"
+def recur_data(dataset_path, dict_label):
+    if os.path.isfile(dataset_path):
+        if ".docx" in dataset_path:
+            return 
+        all_data.append(dataset_path)
+        label_name = dataset_path.split("/")[-2]
+        all_label.append(dict_label[label_name])
+    else:
+        for file_name in os.listdir(dataset_path):
+            new_path = dataset_path + "/" + file_name
+            recur_data(new_path, dict_label)
+    return 
+
+recur_data(dataset_path,dict_label)
+
+# split data to train and test
+train_data, test_data, train_label, test_label = model_selection.train_test_split(all_data, all_label, test_size=configs['split_size'], shuffle=True)
+print("train:", len(train_data), len(train_label), Counter(train_label))
+print("test:", len(test_data), len(test_label), Counter(test_label))
+
+# trrain and test data loader
+class mushroom_Dataset(Dataset):
+    def __init__(self, data,label, transform):
+        self.data = data[:]
+        self.label = label[:]
+        self.transform = transform
+                    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, index):
+        image_path = self.data[index]
+        label = self.label[index]
+        image = Image.open(image_path)
+        image = self.transform(image,return_tensors="pt")
+        return image, label, image_path
+
+image_processor = AutoImageProcessor.from_pretrained("google/bit-50")
+train_dataset = mushroom_Dataset(train_data, train_label, image_processor )
+train_loader = DataLoader(train_dataset, batch_size=configs['batch_size'], 
+                          num_workers=4, shuffle=True, pin_memory=True)
+test_dataset = mushroom_Dataset(test_data, test_label, image_processor )
+test_loader = DataLoader(test_dataset, batch_size=configs['batch_size'], 
+                          num_workers=4, shuffle=True, pin_memory=True)
+
+
+# model class
+class mush_root_model(torch.nn.Module):
+    def __init__(self, num_labels=1):
+        super(mush_root_model, self).__init__()
+        self.model = BitModel.from_pretrained("google/bit-50")
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(2048, num_labels),
+        )
+    
+    def forward(self, input):
+        outputs = self.model(**input).pooler_output
+        #print(outputs.shape)
+        logits = self.classifier(outputs)
+        return logits
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = mush_root_model(num_labels=len(dict_label))
+model.to(device);    
+
+
+# train 
+def loss_fn(outputs, targets):
+    return torch.nn.CrossEntropyLoss()(outputs, targets)
+
+optimizer = AdamW(params =  model.parameters(), lr=configs['LEARNING_RATE'], weight_decay=1e-6)
+
+def validation():
+    model.eval()
+    fin_targets=[]
+    fin_outputs=[]
+    with torch.no_grad():
+        for _, data in enumerate(test_loader, 0):
+            images = data[0].to(device)
+            images['pixel_values'] = torch.squeeze(images['pixel_values'])
+            targets = data[1].to(device, dtype = torch.float)
+            outputs = torch.squeeze(model(images))
+            fin_targets.extend(targets.cpu().detach().numpy().tolist())
+            fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
+    outputs1 = np.argmax(fin_outputs, axis=-1)
+    accuracy = metrics.accuracy_score(fin_targets, outputs1)
+    return accuracy
+
+def train(epoch):
+    model.train()
+    for i in range(epoch):
+        for _,data in enumerate(train_loader, 0):
+            images = data[0].to(device)
+            images['pixel_values'] = torch.squeeze(images['pixel_values'])
+            targets = data[1].to(device, dtype = torch.int64)
+            outputs = torch.squeeze(model(images))
+            loss = loss_fn(outputs, targets)
+            if _%50 == 0:
+                acc = validation()
+                print(f'Epoch: {i}, Loss:  {loss.item()}， val acc: {acc}')
+                model.train()
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+            
+train(configs['EPOCHS'])
+save_path = "/kaggle/working/mush_classifier_20230801.pth"
+torch.save(model.state_dict(), save_path)