Spaces:

kerjaaa
/

trash

No application file

trash / src /train_model.py

Synced repo using 'sync_with_huggingface' Github Action

9076fcb verified almost 2 years ago

4.84 kB

import datasets # type: ignore
import tensorflow as tf # type: ignore
import numpy as np
import random
from PIL import Image
from tensorflow.keras.preprocessing import image  # type: ignore
from tensorflow.keras.applications import ResNet50  # type: ignore
from tensorflow.keras import layers, models  # type: ignore
import wandb # type: ignore
import tensorflow as tf # type: ignore
import sys

# Memeriksa dan mengatur GPU
gpus = tf.config.list_physical_devices('GPU')
if not gpus:
    print("Tidak ada GPU yang terdeteksi. Program ini memerlukan GPU untuk dijalankan.")
    sys.exit(1)  # Menghentikan program dengan kode status 1

try:
    tf.config.experimental.set_memory_growth(gpus[0], True)
    print("Pengaturan memori GPU berhasil.")
except Exception as e:
    print(f"Kesalahan saat mengatur memori GPU: {e}")
    sys.exit(1)  # Menghentikan program jika terjadi kesalahan saat mengatur GPU

# Fungsi untuk memproses dan melakukan augmentasi pada gambar
def preprocess_and_augment_images(dataset, num_images=0, target_size=(224, 224)):
    train_images = []
    train_labels = []

    for i in range(min(num_images, len(dataset['train']))):
        try:
            img = dataset['train'][i]['image'].convert('RGB').resize(target_size)
            label = dataset['train'][i]['label']
            img_array = np.array(img) / 255.0
            
            # Deteksi blur dan augmentasi gambar
            if np.std(img_array) < 0.05:
                img_array *= 2  # Tingkatkan kontras
            
            augment_type = random.choice(['flip', 'rotate', 'brightness'])
            if augment_type == 'flip':
                img_array = np.fliplr(img_array)
            elif augment_type == 'rotate':
                img_array = np.rot90(img_array)
            elif augment_type == 'brightness':
                factor = random.uniform(0.5, 1.5)
                img_array = np.clip(img_array * factor, 0, 1)

            train_images.append(img_array)
            train_labels.append(label)

        except Exception as e:
            print(f"Error processing image {i + 1}: {e}")

    return np.array(train_images), np.array(train_labels)

# Fungsi untuk membuat model ResNet50
def create_resnet50(input_shape, num_classes):
    base_model = ResNet50(weights='imagenet', include_top=False, input_shape=input_shape)
    x = base_model.output
    x = layers.GlobalAveragePooling2D()(x)
    outputs = layers.Dense(num_classes, activation='softmax')(x)

    model = models.Model(inputs=base_model.input, outputs=outputs)

    for layer in base_model.layers:
        layer.trainable = False  # Membekukan lapisan dasar

    return model

# Fungsi untuk melatih dan menyimpan model dengan kondisi tertentu
def train_and_save_model(train_images, train_labels, val_images, val_labels, num_classes, epochs=10):
    model = create_resnet50(input_shape=(224, 224, 3), num_classes=num_classes)

    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    history = model.fit(train_images, train_labels,
                        epochs=epochs,
                        validation_data=(val_images, val_labels))

    # Menyimpan model jika memenuhi syarat akurasi
    if history.history['val_accuracy'][-1] >= 0.80:
        model.save('simpan_resnet50_model.h5')
        print("Model telah disimpan.")
    
    return model

# Fungsi untuk mengevaluasi model
def evaluate_model(model, test_images, test_labels):
    loss, accuracy = model.evaluate(test_images, test_labels)
    print(f"Test Loss: {loss:.4f}, Test Accuracy: {accuracy:.4f}")

# Fungsi utama untuk menjalankan program
def main():
    wandb.login(key="967fafd5558eeeff3ce5681cf55c71633438428d")  # Ganti dengan kunci API Anda
    # Memuat dataset
    ds = datasets.load_dataset("garythung/trashnet")

    num_images = 5024
    train_images, train_labels = preprocess_and_augment_images(ds, num_images=num_images)

    # Memisahkan data menjadi set pelatihan dan validasi
    split_index = int(0.8 * len(train_images))  # 80% untuk pelatihan
    val_images, val_labels = train_images[split_index:], train_labels[split_index:]  # 20% untuk validasi
    train_images, train_labels = train_images[:split_index], train_labels[:split_index]  # 80% untuk pelat ihan

    # Menentukan jumlah kelas
    num_classes = len(set(train_labels))

    # Mengupdate parameter config berdasarkan dataset
    config = {
        "learning_rate": 0.001,
        "epochs": 10,
        "batch_size": 64,
        "architecture": "ResNet50",
        "dataset": ds.info.description,  # Deskripsi dataset
        "num_classes": num_classes  # Jumlah kelas
    }
    
    wandb.init(project='my-tf-trash', config=config)

    model = train_and_save_model(train_images, train_labels, val_images, val_labels, num_classes) 

if __name__ == "__main__":
    main()

	import datasets # type: ignore
	import tensorflow as tf # type: ignore
	import numpy as np
	import random
	from PIL import Image
	from tensorflow.keras.preprocessing import image # type: ignore
	from tensorflow.keras.applications import ResNet50 # type: ignore
	from tensorflow.keras import layers, models # type: ignore
	import wandb # type: ignore
	import tensorflow as tf # type: ignore
	import sys

	# Memeriksa dan mengatur GPU
	gpus = tf.config.list_physical_devices('GPU')
	if not gpus:
	print("Tidak ada GPU yang terdeteksi. Program ini memerlukan GPU untuk dijalankan.")
	sys.exit(1) # Menghentikan program dengan kode status 1

	try:
	tf.config.experimental.set_memory_growth(gpus[0], True)
	print("Pengaturan memori GPU berhasil.")
	except Exception as e:
	print(f"Kesalahan saat mengatur memori GPU: {e}")
	sys.exit(1) # Menghentikan program jika terjadi kesalahan saat mengatur GPU

	# Fungsi untuk memproses dan melakukan augmentasi pada gambar
	def preprocess_and_augment_images(dataset, num_images=0, target_size=(224, 224)):
	train_images = []
	train_labels = []

	for i in range(min(num_images, len(dataset['train']))):
	try:
	img = dataset['train'][i]['image'].convert('RGB').resize(target_size)
	label = dataset['train'][i]['label']
	img_array = np.array(img) / 255.0

	# Deteksi blur dan augmentasi gambar
	if np.std(img_array) < 0.05:
	img_array *= 2 # Tingkatkan kontras

	augment_type = random.choice(['flip', 'rotate', 'brightness'])
	if augment_type == 'flip':
	img_array = np.fliplr(img_array)
	elif augment_type == 'rotate':
	img_array = np.rot90(img_array)
	elif augment_type == 'brightness':
	factor = random.uniform(0.5, 1.5)
	img_array = np.clip(img_array * factor, 0, 1)

	train_images.append(img_array)
	train_labels.append(label)

	except Exception as e:
	print(f"Error processing image {i + 1}: {e}")

	return np.array(train_images), np.array(train_labels)

	# Fungsi untuk membuat model ResNet50
	def create_resnet50(input_shape, num_classes):
	base_model = ResNet50(weights='imagenet', include_top=False, input_shape=input_shape)
	x = base_model.output
	x = layers.GlobalAveragePooling2D()(x)
	outputs = layers.Dense(num_classes, activation='softmax')(x)

	model = models.Model(inputs=base_model.input, outputs=outputs)

	for layer in base_model.layers:
	layer.trainable = False # Membekukan lapisan dasar

	return model

	# Fungsi untuk melatih dan menyimpan model dengan kondisi tertentu
	def train_and_save_model(train_images, train_labels, val_images, val_labels, num_classes, epochs=10):
	model = create_resnet50(input_shape=(224, 224, 3), num_classes=num_classes)

	model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

	history = model.fit(train_images, train_labels,
	epochs=epochs,
	validation_data=(val_images, val_labels))

	# Menyimpan model jika memenuhi syarat akurasi
	if history.history['val_accuracy'][-1] >= 0.80:
	model.save('simpan_resnet50_model.h5')
	print("Model telah disimpan.")

	return model

	# Fungsi untuk mengevaluasi model
	def evaluate_model(model, test_images, test_labels):
	loss, accuracy = model.evaluate(test_images, test_labels)
	print(f"Test Loss: {loss:.4f}, Test Accuracy: {accuracy:.4f}")

	# Fungsi utama untuk menjalankan program
	def main():
	wandb.login(key="967fafd5558eeeff3ce5681cf55c71633438428d") # Ganti dengan kunci API Anda
	# Memuat dataset
	ds = datasets.load_dataset("garythung/trashnet")

	num_images = 5024
	train_images, train_labels = preprocess_and_augment_images(ds, num_images=num_images)

	# Memisahkan data menjadi set pelatihan dan validasi
	split_index = int(0.8 * len(train_images)) # 80% untuk pelatihan
	val_images, val_labels = train_images[split_index:], train_labels[split_index:] # 20% untuk validasi
	train_images, train_labels = train_images[:split_index], train_labels[:split_index] # 80% untuk pelat ihan

	# Menentukan jumlah kelas
	num_classes = len(set(train_labels))

	# Mengupdate parameter config berdasarkan dataset
	config = {
	"learning_rate": 0.001,
	"epochs": 10,
	"batch_size": 64,
	"architecture": "ResNet50",
	"dataset": ds.info.description, # Deskripsi dataset
	"num_classes": num_classes # Jumlah kelas
	}

	wandb.init(project='my-tf-trash', config=config)

	model = train_and_save_model(train_images, train_labels, val_images, val_labels, num_classes)

	if __name__ == "__main__":
	main()