Property	Value
Base Model	`convnext_base.fb_in22k_ft_in1k_384`
Pretraining	ImageNet-22k then ImageNet-1k (86.8% top-1)
Input Size	384 x 384 px
Parameters	~88M
Fine-tuning	Full LLRD + EMA

Classes

Label	Code	Material	Recyclable
PET	#1	Polyethylene Terephthalate	Yes
HDPE	#2	High-Density Polyethylene	Yes
LDPE	#4	Low-Density Polyethylene	Yes
PP	#5	Polypropylene	Yes
PS	#6	Polystyrene	No
Other	#7	Mixed / Unknown	Depends

Quick Start

python

import torch, timm, numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from huggingface_hub import hf_hub_download
from PIL import Image

weights = hf_hub_download('Vansh180/plastic-waste-classifier', 'plastic_classifier_best.pt')
model = timm.create_model('convnext_base.fb_in22k_ft_in1k_384', pretrained=False, num_classes=6)
model.load_state_dict(torch.load(weights, map_location='cpu'))
model.eval()

transform = A.Compose([
    A.Resize(416, 416), A.CenterCrop(384, 384),
    A.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]), ToTensorV2(),
])
CLASS_NAMES = ['HDPE','LDPE','Other','PET','PP','PS']
img = np.array(Image.open('plastic.jpg').convert('RGB'))
t = transform(image=img)['image'].unsqueeze(0)
with torch.no_grad():
    probs = torch.softmax(model(t), dim=1)[0]
print(CLASS_NAMES[probs.argmax()], f'{probs.max():.1%}')

License

Apache 2.0