File size: 5,239 Bytes
caaad53 968fed8 caaad53 91ca776 caaad53 573f012 caaad53 63a54d3 caaad53 968fed8 caaad53 968fed8 caaad53 968fed8 caaad53 302855c 32b3004 caaad53 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import os
import shutil
import subprocess
import zipfile
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.optim import lr_scheduler
import subprocess
import zipfile
from PIL import Image
import gradio as gr
# Setup Kaggle API
kaggle_dir = os.path.expanduser("~/.kaggle")
if not os.path.exists(kaggle_dir):
os.makedirs(kaggle_dir)
# Copy the kaggle.json file to the ~/.kaggle directory
kaggle_json_path = "kaggle.json"
kaggle_dest_path = os.path.join(kaggle_dir, "kaggle.json")
if not os.path.exists(kaggle_dest_path):
shutil.copy(kaggle_json_path, kaggle_dest_path)
os.chmod(kaggle_dest_path, 0o600)
print("Kaggle API key copied and permissions set.")
else:
print("Kaggle API key already exists.")
# Download the dataset from Kaggle using Kaggle CLI
dataset_name = "mostafaabla/garbage-classification"
print(f"Downloading the dataset: {dataset_name}")
download_command = f"kaggle datasets download -d {dataset_name}"
# Run the download command
subprocess.run(download_command, shell=True)
# Unzip the downloaded dataset
dataset_zip = "garbage-classification.zip"
extracted_folder = "./garbage-classification"
# Check if the zip file exists
if os.path.exists(dataset_zip):
if not os.path.exists(extracted_folder):
with zipfile.ZipFile(dataset_zip, 'r') as zip_ref:
zip_ref.extractall(extracted_folder)
print("Dataset unzipped successfully!")
else:
print("Dataset already unzipped.")
else:
print(f"Dataset zip file '{dataset_zip}' not found.")
# Model training and testing in separate directory at ipynb file (Copy of ai-portfolio Kendrick.ipynb)
from PIL import Image
import gradio as gr
# Load model
def load_model():
model = models.resnet50(weights='DEFAULT') # Using default weights for initialization
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 12) # Adjust to the number of classes you have
# Load the state dict
model.load_state_dict(torch.load('resnet50_garbage_classificationv1.2.pth', map_location=torch.device('cpu')))
model.eval() # Set to evaluation mode
return model
model = load_model()
# Define image transformations
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# Class names
class_names = ['battery', 'biological', 'brown-glass', 'cardboard',
'clothes', 'green-glass', 'metal', 'paper',
'plastic', 'shoes', 'trash', 'white-glass']
# Define bin colors for each class
bin_colors = {
'battery': 'Merah (Red)', # Limbah berbahaya (B3)
'biological': 'Hijau (Green)', # Limbah organik
'brown-glass': 'Kuning (Yellow or trash banks / recycling centers)', # Gelas berwarna coklat (anorganik/daur ulang)
'cardboard': 'Biru (Blue)', # Kertas (daur ulang)
'clothes': 'Kuning atau Bank Sampah (Yellow or trash banks / recycling centers)', # Pakaian (dimasukkan sebagai daur ulang)
'green-glass': 'Kuning (Yellow)', # Gelas berwarna hijau (anorganik/daur ulang)
'metal': 'Kuning (Yellow)', # Logam (anorganik/daur ulang)
'paper': 'Biru (Blue)', # Kertas (daur ulang)
'plastic': 'Kuning (Yellow)', # Plastik (anorganik/daur ulang)
'shoes': 'Kuning atau Bank Sampah (Yellow or trash banks / recycling centers)', # Sepatu (dimasukkan sebagai daur ulang)
'trash': 'Abu-abu (Gray)', # Limbah umum
'white-glass': 'Kuning (Yellow or trash banks / recycling centers)' # Gelas berwarna putih (anorganik/daur ulang)
}
# Define the prediction function
def predict(image):
image = Image.fromarray(image) # Convert numpy array to PIL Image
image = transform(image) # Apply transformations
image = image.unsqueeze(0) # Add batch dimension
with torch.no_grad():
outputs = model(image)
_, predicted = torch.max(outputs, 1)
class_name = class_names[predicted.item()] # Return predicted class name
bin_color = bin_colors[class_name] # Get the corresponding bin color
return class_name, bin_color # Return both class name and bin color
# Make Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="numpy", label="Unggah Gambar"),
outputs=[
gr.Textbox(label="Jenis Sampah"),
gr.Textbox(label="Tong Sampah yang Sesuai") # 2 output with label
],
title="Klasifikasi Sampah dengan ResNet50 v1.2",
description="Unggah gambar sampah, dan model kami akan mengklasifikasikannya ke dalam salah satu dari 12 kategori bersama dengan warna tempat sampah yang sesuai. "
"<strong>Model ini bisa memprediksi jenis sampah dari ke-12 jenis berikut:</strong> Baterai, Sampah organik, Gelas Kaca Coklat, "
"Kardus, Pakaian, Gelas Kaca Hijau, Metal, Kertas, Plastik, Sepatu/sandal, Popok/pampers, Gelas Kaca bening. "
"<strong>Note: Untuk masker dan pampers dikategorikan sebagai trash</strong>"
)
iface.launch(share=True)
|