Spaces:
Sleeping
Sleeping
File size: 5,661 Bytes
52dd1ca |
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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import os
import json
from typing import Tuple
import numpy as np
from tqdm import tqdm
import torch
from torchvision import datasets, transforms
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report
import joblib
def get_datasets(data_root: str, image_size: int = 64) -> Tuple[torch.utils.data.Dataset,
torch.utils.data.Dataset,
dict]:
"""
Load Oxford-IIIT Pet train/test splits with simple transforms.
Returns:
train_dataset, test_dataset, class_to_idx
"""
# Simple transform: resize -> grayscale -> tensor in [0,1]
transform = transforms.Compose([
transforms.Resize((image_size, image_size)),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor(), # (1, H, W), float32 in [0,1]
])
train_dataset = datasets.OxfordIIITPet(
root=data_root,
split="trainval",
target_types="category",
transform=transform,
download=True, # downloads to root/oxford-iiit-pet if not present
)
test_dataset = datasets.OxfordIIITPet(
root=data_root,
split="test",
target_types="category",
transform=transform,
download=True,
)
# class_to_idx mapping
# Many torchvision datasets expose this attribute
class_to_idx = train_dataset.class_to_idx
return train_dataset, test_dataset, class_to_idx
def dataset_to_numpy(dataset: torch.utils.data.Dataset) -> Tuple[np.ndarray, np.ndarray]:
"""
Convert a torchvision dataset (with tensor images) to numpy arrays
suitable for scikit-learn.
X: (N, D) flattened grayscale pixels
y: (N,) int labels
"""
X_list = []
y_list = []
for img, label in tqdm(dataset, desc="Converting to numpy"):
# img: torch.Tensor, shape (1, H, W)
arr = img.numpy() # (1, H, W)
arr = arr.reshape(-1) # flatten to (D,)
X_list.append(arr)
y_list.append(label)
X = np.stack(X_list, axis=0).astype(np.float32) # (N, D)
y = np.array(y_list, dtype=np.int64) # (N,)
return X, y
def save_labels(class_to_idx: dict, labels_path: str):
"""
Save labels as id -> class_name in a JSON file for inference/UI.
"""
# Invert mapping: idx -> class_name
idx_to_class = {idx: cls_name for cls_name, idx in class_to_idx.items()}
os.makedirs(os.path.dirname(labels_path), exist_ok=True)
with open(labels_path, "w") as f:
json.dump(idx_to_class, f, indent=2)
print(f"[INFO] Saved labels to {labels_path}")
def train_logistic_regression(X_train: np.ndarray, y_train: np.ndarray) -> LogisticRegression:
"""
Train multinomial Logistic Regression on given features.
We use 'saga' because it supports multinomial loss and L1/L2,
and works decently with high-dimensional sparse-ish data.
"""
num_classes = len(np.unique(y_train))
print(f"[INFO] Training Logistic Regression on {X_train.shape[0]} samples, "
f"{X_train.shape[1]} features, {num_classes} classes")
clf = LogisticRegression(
penalty="l2",
C=1.0,
solver="saga",
multi_class="multinomial",
max_iter=1000,
n_jobs=-1,
verbose=1,
)
clf.fit(X_train, y_train)
return clf
def evaluate_model(clf: LogisticRegression, X: np.ndarray, y: np.ndarray, split_name: str):
"""
Print accuracy and basic classification report for a given split.
"""
y_pred = clf.predict(X)
acc = accuracy_score(y, y_pred)
print(f"\n[{split_name}] Accuracy: {acc * 100:.2f}%")
print(f"[{split_name}] Classification report (macro avg at bottom):")
print(classification_report(y, y_pred, digits=3))
def main():
# -------- configs (tweak paths as needed) --------
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
data_root = os.path.join(project_root, "data")
checkpoints_dir = os.path.join(project_root, "checkpoints")
configs_dir = os.path.join(project_root, "configs")
os.makedirs(checkpoints_dir, exist_ok=True)
os.makedirs(configs_dir, exist_ok=True)
labels_path = os.path.join(configs_dir, "labels.json")
model_path = os.path.join(checkpoints_dir, "lr_model.joblib")
image_size = 64 # 64x64 grayscale baseline
# ------------------------------------------------
print("[INFO] Loading datasets...")
train_dataset, test_dataset, class_to_idx = get_datasets(data_root, image_size=image_size)
print(f"[INFO] Train samples: {len(train_dataset)}, Test samples: {len(test_dataset)}")
print(f"[INFO] Number of classes: {len(class_to_idx)}")
print("[INFO] Converting train split to numpy...")
X_train, y_train = dataset_to_numpy(train_dataset)
print("[INFO] Converting test split to numpy...")
X_test, y_test = dataset_to_numpy(test_dataset)
# Save label mapping for later inference
save_labels(class_to_idx, labels_path)
# Train LR
clf = train_logistic_regression(X_train, y_train)
# Evaluate
evaluate_model(clf, X_train, y_train, split_name="Train")
evaluate_model(clf, X_test, y_test, split_name="Test")
# Save model
joblib.dump(clf, model_path)
print(f"[INFO] Saved Logistic Regression model to {model_path}")
if __name__ == "__main__":
main()
|