Spaces:
Runtime error
Runtime error
File size: 5,254 Bytes
dec6452 | 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 | # Import necessary libraries
import torch
import torchvision.models as models
from torchvision import transforms
from PIL import Image
import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
import gradio as gr
# Check if GPU is available and use it
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Define transformation for data preprocessing and augmentation
transform = transforms.Compose([
transforms.Resize((224, 224)), # Resize images to 224x224
transforms.ToTensor(), # Convert image to tensor
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize image
])
augment_transform = transforms.Compose([
transforms.RandomHorizontalFlip(), # Randomly flip images horizontally
transforms.RandomRotation(10), # Randomly rotate images by 10 degrees
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2), # Randomly change image brightness, contrast, etc.
transforms.Resize((224, 224)), # Resize images to 224x224
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize image
])
# Function to load and preprocess images
def load_and_process_images(image_folder, transform, augment_transform):
processed_images = []
labels = []
for label_idx, label_name in enumerate(os.listdir(image_folder)):
label_folder = os.path.join(image_folder, label_name)
if os.path.isdir(label_folder):
for filename in os.listdir(label_folder):
if filename.endswith('.jpg') or filename.endswith('.png'):
image_path = os.path.join(label_folder, filename)
image = Image.open(image_path).convert('RGB')
processed_images.append(transform(image))
labels.append(label_idx)
# Apply augmentation and add augmented image
augmented_image = augment_transform(image)
processed_images.append(augmented_image)
labels.append(label_idx)
return torch.stack(processed_images), labels
# Load pre-trained ResNet-50 model
model = models.resnet50(pretrained=True)
model = torch.nn.Sequential(*(list(model.children())[:-1])) # Remove the final fully connected layer
model = model.to(device)
model.eval()
# Function to extract features using ResNet-50
def extract_features(images, model, batch_size=32):
model.eval()
features = []
with torch.no_grad():
for i in range(0, len(images), batch_size):
batch_images = images[i:i + batch_size].to(device)
batch_features = model(batch_images).squeeze().cpu().numpy()
features.append(batch_features)
return np.vstack(features)
# Function for K-means clustering
def perform_kmeans(features, num_clusters):
kmeans = KMeans(n_clusters=num_clusters, random_state=42)
predicted_labels = kmeans.fit_predict(features)
return kmeans, predicted_labels
# Specify image folder containing dataset
image_folder = '/content/drive/MyDrive/Colab Notebooks/Datasets/defects_kaggle'
# Load and preprocess images
processed_images, labels = load_and_process_images(image_folder, transform, augment_transform)
processed_images = processed_images.to(device)
# Extract features using ResNet-50
features = extract_features(processed_images, model)
# Perform K-means clustering on extracted features
num_clusters = len(os.listdir(image_folder)) # Number of clusters based on folder names
kmeans, predicted_labels = perform_kmeans(features, num_clusters)
# Reduce dimensionality for visualization
pca = PCA(n_components=50)
pca_features = pca.fit_transform(features)
tsne = TSNE(n_components=2, random_state=42)
tsne_features = tsne.fit_transform(pca_features)
# Plot clusters
plt.figure(figsize=(10, 8))
scatter = plt.scatter(tsne_features[:, 0], tsne_features[:, 1], c=predicted_labels, cmap='viridis')
plt.title('Clusters of Images')
plt.xlabel('TSNE Feature 1')
plt.ylabel('TSNE Feature 2')
plt.legend(handles=scatter.legend_elements()[0], labels=os.listdir(image_folder))
plt.colorbar(scatter)
plt.show()
# Function to predict the cluster for a new image
def predict_cluster(image):
image = Image.fromarray(image).convert('RGB')
image = transform(image).unsqueeze(0).to(device)
features = extract_features(image, model)
cluster = kmeans.predict(features)
return f"Predicted Cluster: {cluster[0]}"
# Create a mapping from cluster numbers to directory names
cluster_names = {i: name for i, name in enumerate(os.listdir(image_folder))}
def predict_cluster(image):
image = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
feature = model(image).squeeze().cpu().numpy().reshape(1, -1)
cluster = kmeans.predict(feature)[0]
cluster_name = cluster_names[cluster]
return f"Cluster: {cluster}, Name: {cluster_name}"
gr.Interface(
fn=predict_cluster,
inputs=gr.Image(type="pil"),
outputs="text",
title="Image Cluster Prediction",
description="Upload an image to predict its cluster."
).launch()
|