File size: 3,336 Bytes
028930f 5d55f41 028930f 5d55f41 028930f 5d55f41 028930f |
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 |
import os
import numpy as np
import cv2
from sklearn.model_selection import train_test_split
from transformers import ViTForImageClassification, ViTFeatureExtractor
from transformers import Trainer, TrainingArguments
from datasets import load_dataset, Dataset
# Function to load images from specified folders
def load_images_from_folders(folders, label):
images = []
labels = []
for folder in folders:
for filename in os.listdir(folder):
if filename.lower().endswith(('.png', '.jpg', '.jpeg')): # Check for valid image extensions
img = cv2.imread(os.path.join(folder, filename), cv2.IMREAD_GRAYSCALE) # Read image as grayscale
if img is not None:
img = cv2.resize(img, (224, 224)) # Resize to 224x224 pixels
img = img.astype(np.float32) # Ensure the image is in float32 format
img /= 255.0 # Normalize to [0, 1]
images.append(img)
labels.append(label)
else:
print(f"Failed to load image: {filename}")
return images, labels
# Load normal and pneumonia images
normal_folders = [
os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'NORMAL'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'NORMAL'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'NORMAL'),
]
pneumonia_folders = [
os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'PNEUMONIA'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'PNEUMONIA'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'PNEUMONIA'),
]
normal_images, normal_labels = load_images_from_folders(normal_folders, 0)
pneumonia_images, pneumonia_labels = load_images_from_folders(pneumonia_folders, 1)
# Combine images and labels
images = normal_images + pneumonia_images
labels = normal_labels + pneumonia_labels
# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)
# Convert the dataset to a Hugging Face Dataset
train_dataset = Dataset.from_dict({"image": X_train, "label": y_train})
test_dataset = Dataset.from_dict({"image": X_test, "label": y_test})
# Load feature extractor and model
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=2)
# Preprocess the dataset
def preprocess_function(examples):
return feature_extractor(images=examples['image'], return_tensors="pt")
train_dataset = train_dataset.map(preprocess_function, batched=True)
test_dataset = test_dataset.map(preprocess_function, batched=True)
# Training arguments
training_args = TrainingArguments(
output_dir='./results',
evaluation_strategy='epoch',
learning_rate=2e-5,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=10,
weight_decay=0.01,
)
# Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=test_dataset,
)
# Train the model
trainer.train()
# Save the model
model.save_pretrained('./pneumonia_model_final')
print("Model saved as './pneumonia_model_final'")
|