Upload 6 files

README.md

@@ -0,0 +1,3 @@
+Training of https://github.com/yuechen-yang/garbage-classification using this dataset: https://www.kaggle.com/datasets/mostafaabla/garbage-classification
+
+

config.json

@@ -0,0 +1,52 @@
+{
+  "_name_or_path": "google/vit-base-patch16-224-in21k",
+  "architectures": [
+    "ViTForImageClassification"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "encoder_stride": 16,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "battery",
+    "1": "biological",
+    "10": "trash",
+    "11": "white-glass",
+    "2": "brown-glass",
+    "3": "cardboard",
+    "4": "clothes",
+    "5": "green-glass",
+    "6": "metal",
+    "7": "paper",
+    "8": "plastic",
+    "9": "shoes"
+  },
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "battery": "0",
+    "biological": "1",
+    "brown-glass": "2",
+    "cardboard": "3",
+    "clothes": "4",
+    "green-glass": "5",
+    "metal": "6",
+    "paper": "7",
+    "plastic": "8",
+    "shoes": "9",
+    "trash": "10",
+    "white-glass": "11"
+  },
+  "layer_norm_eps": 1e-12,
+  "model_type": "vit",
+  "num_attention_heads": 12,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "patch_size": 16,
+  "problem_type": "single_label_classification",
+  "qkv_bias": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.1"
+}

inference.py

@@ -0,0 +1,25 @@
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
+
+model_name = "saved_model"
+
+model = ViTForImageClassification.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+
+model.eval()
+
+image_path = '/path/'
+image = Image.open(image_path).convert('RGB')
+
+inputs = feature_extractor(images=image, return_tensors="pt")
+
+with torch.no_grad():
+    outputs = model(**inputs)
+    logits = outputs.logits
+
+predicted_class_idx = logits.argmax(-1).item()
+
+classes = model.config.id2label
+
+print(f"Predicted class: {classes[predicted_class_idx]}")

model.safetensors

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

preprocessor_config.json

@@ -0,0 +1,22 @@
+{
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "ViTFeatureExtractor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 224,
+    "width": 224
+  }
+}

train.py

@@ -0,0 +1,112 @@
+import torch
+import math
+from torchvision.datasets import ImageFolder
+from torch.utils.data import DataLoader
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+import kagglehub
+from torch.optim import AdamW
+from transformers import get_scheduler
+from tqdm.auto import tqdm
+
+
+gpu_available = torch.cuda.is_available()
+print("GPU available:", gpu_available)
+
+if gpu_available:
+    print("GPU:", torch.cuda.get_device_name(0))
+    print("GPU count:", torch.cuda.device_count())
+    print("#memory avail:", torch.cuda.get_device_properties(0).total_memory / (1024 ** 3), "GB")
+else:
+    print("No GPU available.")
+
+
+kaggle_path = kagglehub.dataset_download("mostafaabla/garbage-classification")
+
+folder_root_path = kaggle_path+'/garbage_classification'
+print("Path to dataset files:", kaggle_path)
+
+ds = ImageFolder(folder_root_path)
+indices = torch.randperm(len(ds)).tolist()
+n_val = math.floor(len(indices) * .20)
+train_ds = torch.utils.data.Subset(ds, indices[:-n_val])
+val_ds = torch.utils.data.Subset(ds, indices[-n_val:])
+
+print(ds.classes)
+
+label2id = {}
+id2label = {}
+
+for i, class_name in enumerate(ds.classes):
+    label2id[class_name] = str(i)
+    id2label[str(i)] = class_name
+
+
+class ImageClassificationCollator:
+    def __init__(self, feature_extractor):
+        self.feature_extractor = feature_extractor
+
+    def __call__(self, batch):
+        encodings = self.feature_extractor([x[0] for x in batch], return_tensors='pt')
+        encodings['labels'] = torch.tensor([x[1] for x in batch], dtype=torch.long)
+        return encodings
+
+
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+
+model = ViTForImageClassification.from_pretrained(
+    'google/vit-base-patch16-224-in21k',
+    num_labels=len(label2id),
+    label2id=label2id,
+    id2label=id2label
+)
+collator = ImageClassificationCollator(feature_extractor)
+train_loader = DataLoader(train_ds, batch_size=16, collate_fn=collator, num_workers=2, shuffle=True)
+val_loader = DataLoader(val_ds, batch_size=16, collate_fn=collator, num_workers=2)
+
+optimizer = AdamW(model.parameters(), lr=5e-5)
+
+num_epochs = 10
+num_training_steps = num_epochs * len(train_loader)
+lr_scheduler = get_scheduler(
+    name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
+)
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+model.to(device)
+
+
+progress_bar = tqdm(range(num_training_steps))
+
+model.train()
+for epoch in range(num_epochs):
+    for batch in train_loader:
+        batch = {k: v.to(device) for k, v in batch.items()}
+        outputs = model(**batch)
+        loss = outputs.loss
+        loss.backward()
+
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad()
+        progress_bar.update(1)
+
+
+model.eval()
+for batch in val_loader:
+    batch = {k: v.to(device) for k, v in batch.items()}
+    with torch.no_grad():
+        outputs = model(**batch)
+
+    logits = outputs.logits
+    predictions = torch.argmax(logits, dim=-1)
+ 
+import os
+
+save_directory = "./saved_model"
+if not os.path.exists(save_directory):
+    os.makedirs(save_directory)
+
+model.save_pretrained(save_directory)
+feature_extractor.save_pretrained(save_directory)
+
+print(f"Model saved: {save_directory}")