Function main

ML/src/python/neuralforge/cli/gui.py:4–486

()

import os

def main():
    try:
        from PyQt6.QtWidgets import QApplication
    except ImportError:
        print("Error: PyQt6 not installed")
        print("Install with: pip install neuralforge[gui]")
        print("Or: pip install PyQt6")
        sys.exit(1)
    
    current_dir = os.path.dirname(os.path.abspath(__file__))
    root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_dir))))
    
    sys.path.insert(0, root_dir)
    
    from PyQt6.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                                 QPushButton, QLabel, QLineEdit, QFileDialog, 
                                 QProgressBar, QTextEdit, QGroupBox)
    from PyQt6.QtCore import Qt, QThread, pyqtSignal
    from PyQt6.QtGui import QPixmap, QFont
    
    import torch
    import torch.nn.functional as F
    from torchvision import transforms
    from PIL import Image
    
    from neuralforge.data.datasets import get_dataset, get_num_classes
    from neuralforge.models.resnet import ResNet18
    
    class PredictionThread(QThread):
        finished = pyqtSignal(list, list, str)
        error = pyqtSignal(str)
        
        def __init__(self, model, image_path, classes, device):
            super().__init__()
            self.model = model
            self.image_path = image_path
            self.classes = classes
            self.device = device
        
        def run(self):
            try:
                image = Image.open(self.image_path).convert('RGB')
                
                transform = transforms.Compose([
                    transforms.Resize(256),
                    transforms.CenterCrop(224),
                    transforms.ToTensor(),
                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                ])
                
                image_tensor = transform(image).unsqueeze(0).to(self.device)
                
                with torch.no_grad():
                    outputs = self.model(image_tensor)
                    probabilities = F.softmax(outputs, dim=1)
                    
                    top5_prob, top5_idx = torch.topk(probabilities, min(5, len(self.classes)), dim=1)