refactor: переделана логика

.gitignore

@@ -16,6 +16,6 @@ uv.lock
 .vscode
 test.drawio
 test.py
-models/*
+models/
 Fashion-MNIST/
 MNIST/

dataset_downloader.py

@@ -0,0 +1,50 @@
+import os
+import argparse
+
+import pandas as pd
+from datasets import load_dataset
+
+
+def export_mnist_splits(root_dir: str, dataset: str):
+    ds = load_dataset(dataset)
+
+    img_dir = os.path.join(root_dir, "img")
+    os.makedirs(img_dir, exist_ok=True)
+
+    def save_split(split_name: str):
+        if split_name not in ds:
+            print(f"Сплит '{split_name}' не найден в датасете {dataset}, пропускаю")
+            return
+
+        split = ds[split_name]
+        rows = []
+
+        for idx, example in enumerate(split):
+            img = example["image"]
+            label = example["label"]
+
+            filename = f"{split_name}_{idx:05d}.png"
+            rel_path = f"img/{filename}"
+            abs_path = os.path.join(img_dir, filename)
+
+            img.save(abs_path)
+            rows.append({"path": rel_path, "label": label})
+
+        csv_path = os.path.join(root_dir, f"{split_name}.csv")
+        df = pd.DataFrame(rows)
+        df.to_csv(csv_path, index=False)
+
+        print(f"{split_name}.csv сохранён в {csv_path}, изображений: {len(split)}")
+
+    save_split("train")
+    save_split("test")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", type=str, required=True)
+    parser.add_argument("-d", "--dataset", type=str, required=True)
+
+    args = parser.parse_args()
+
+    export_mnist_splits(args.folder, args.dataset)

mnist.py

@@ -1,12 +1,44 @@
-import argparse
-
-import numpy as np
 import pandas as pd
+import argparse
 
+import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import DataLoader, TensorDataset
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as transforms
+from PIL import Image
+
+
+class ImagePathDataset(Dataset):
+    def __init__(self, csv_path):
+        self.base_dir = os.path.dirname(csv_path)
+
+        df = pd.read_csv(csv_path)
+
+        self.paths = df["path"].tolist()
+        self.labels = df["label"].astype(int).tolist()
+
+        self.transform = transforms.Compose(
+            [
+                transforms.Grayscale(num_output_channels=1),
+                transforms.Resize((28, 28)),
+                transforms.ToTensor(),
+            ]
+        )
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        rel_path = self.paths[idx]
+        full_path = os.path.join(self.base_dir, rel_path)
+
+        img = Image.open(full_path).convert("L")
+        img = self.transform(img)
+
+        label = self.labels[idx]
+        return img, label, rel_path
 
 
 class ModelCNN(nn.Module):
@@ -15,10 +47,10 @@ class ModelCNN(nn.Module):
     INPUT (1x28x28) ->
       [CONV -> RELU -> CONV -> RELU -> POOL] * 3 ->
       [FC -> RELU] * 2 ->
-      FC (num_classes)
+      FC (10)
     """
 
-    def __init__(self, num_classes=10):
+    def __init__(self):
         super(ModelCNN, self).__init__()
 
         self.features = nn.Sequential(
@@ -47,7 +79,7 @@ class ModelCNN(nn.Module):
             nn.ReLU(inplace=True),
             nn.Linear(256, 128),
             nn.ReLU(inplace=True),
-            nn.Linear(128, num_classes),
+            nn.Linear(128, 10),
         )
 
     def forward(self, x):
@@ -61,24 +93,16 @@ def train_mode(args):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Устройство: {device}")
 
-    df = pd.read_csv(args.input)
-
-    labels = df["label"].astype(np.int64).values
-    pixels = df.drop(columns=["label"]).values.astype(np.float32) / 255.0
-
-    images = torch.from_numpy(pixels.reshape(-1, 1, 28, 28))
-    labels = torch.from_numpy(labels)
-
-    dataset = TensorDataset(images, labels)
+    dataset = ImagePathDataset(args.dataset)
     dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
 
-    model = ModelCNN(num_classes=args.num_classes).to(device)
+    model = ModelCNN().to(device)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=args.lr)
 
     model.train()
     for epoch in range(args.epochs):
-        for i, (images, labels) in enumerate(dataloader):
+        for i, (images, labels, _) in enumerate(dataloader):
             images = images.to(device)
             labels = labels.to(device)
 
@@ -94,50 +118,40 @@ def train_mode(args):
                     f"Epoch [{epoch+1}/{args.epochs}], Step [{i+1}/{len(dataloader)}], Loss: {loss.item():.4f}"
                 )
 
-    checkpoint = {
-        "state_dict": model.state_dict(),
-        "num_classes": args.num_classes,
-    }
-    torch.save(checkpoint, args.model)
+    torch.save(model.state_dict(), args.model)
 
 
 def inference_mode(args):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Устройство: {device}")
 
-    checkpoint = torch.load(args.model, map_location=device)
-    num_classes = checkpoint.get("num_classes", 10)
+    state_dict = torch.load(args.model, map_location=device)
 
-    model = ModelCNN(num_classes=num_classes).to(device)
-    model.load_state_dict(checkpoint["state_dict"])
+    model = ModelCNN().to(device)
+    model.load_state_dict(state_dict)
     model.eval()
 
-    df_test = pd.read_csv(args.input)
-
-    has_label = "label" in df_test.columns
-    if has_label:
-        pixels = df_test.drop(columns=["label"]).values
-    else:
-        pixels = df_test.values
-
-    pixels = pixels.astype(np.float32) / 255.0
-    images = torch.from_numpy(pixels.reshape(-1, 1, 28, 28))
-
-    dataset = TensorDataset(images)
+    dataset = ImagePathDataset(args.input)
     dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
 
     all_preds = []
+    all_paths = []
 
     with torch.no_grad():
-        for (batch_images,) in dataloader:
-            batch_images = batch_images.to(device)
-            outputs = model(batch_images)
+        for images, _, pathes in dataloader:
+            images = images.to(device)
+            outputs = model(images)
             _, preds = torch.max(outputs, 1)
-            all_preds.extend(preds.cpu().numpy().tolist())
-
-    df_pred = df_test.copy()
-    df_pred["label"] = all_preds
 
+            all_preds.extend(preds.cpu().numpy().tolist())
+            all_paths.extend(pathes)
+
+    df_pred = pd.DataFrame(
+        {
+            "path": all_paths,
+            "label": all_preds,
+        }
+    )
     df_pred.to_csv(args.output, index=False)
 
 
@@ -145,6 +159,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
 
     parser.add_argument("--mode", choices=["train", "inference"], required=True)
+    parser.add_argument("--dataset", type=str)
     parser.add_argument("--input", type=str)
     parser.add_argument("--output", type=str)
     parser.add_argument("--model", type=str, required=True)
@@ -152,13 +167,12 @@ def parse_args():
     parser.add_argument("--epochs", type=int, default=5)
     parser.add_argument("--batch-size", type=int, default=64)
     parser.add_argument("--lr", type=float, default=0.001)
-    parser.add_argument("--num-classes", type=int, default=10)
 
     args = parser.parse_args()
 
     if args.mode == "train":
-        if args.input is None:
-            parser.error("--input обязателен в режиме train")
+        if args.dataset is None:
+            parser.error("--dataset обязателен в режиме train")
     elif args.mode == "inference":
         if args.input is None or args.output is None:
             parser.error("--input и --output обязательны в режиме inference")