Release/App.py - master - stupin/AutoDriver

Release/App.py
56 строк2 KB
Stupin
init
25 июн 2025, 19:30
25 июн 2025, 19:30a4828f6
from torchvision import io
from PIL import Image
from torch import nn
import torch

class Guide(nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.conv1 = nn.Sequential(nn.Conv2d(3, 6, kernel_size=5, padding=2),
                                   nn.ReLU(inplace=True),
                                   nn.MaxPool2d(2, 2))
        self.conv2 = nn.Sequential(nn.Conv2d(6, 16, kernel_size=5, padding=2),
                                   nn.ReLU(inplace=True),
                                   nn.MaxPool2d(2, 2))
        self.lin1 = nn.Sequential(nn.Flatten(),
                                  nn.Linear(16 * 32 * 32, 9))

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.lin1(out)
        return out

def questguide() -> None:
    DEVICE = torch.device(f'cuda:{0}' if torch.cuda.is_available() else 'cpu')
    model = Guide().to(DEVICE)
    PATH = 'mapper_model.h5'
    model.load_state_dict(torch.load(PATH))
    model.eval()
    for i in range(4214, 4224):
        img = io.read_image(f'dataset/quests/{i}.png', mode=io.ImageReadMode.RGB).to(DEVICE).unsqueeze(dim=0).float() / 256
        out = model(img)
        minimap = Image.open(f'dataset/quests/{i}.png')
        minimap.putpixel((int(out[0][1] * 128), int(out[0][2] * 128)), (255, 0, 0))
        minimap.putpixel((int(out[0][3] * 128), int(out[0][4] * 128)), (255, 0, 0))
        minimap.putpixel((int(out[0][5] * 128), int(out[0][6] * 128)), (255, 0, 0))
        minimap.putpixel((int(out[0][7] * 128), int(out[0][8] * 128)), (255, 0, 0))
        minimap.save(f'temp/{i}-{int(out[0][0] * 360)}.png')

def tracker() -> None:
    pass

def main() -> None:
    print('Welcome to ArcheAge DL Helper!\nAvailable mods:\n1 - Auto quest walker\n2 - Eggs deliver')
    mode = int(input('Choose mode - '))
    match mode:
        case 1:
            questguide()
        case 2:
            tracker()
        case _:
            print('Uncorrect mode! Exiting...')

if __name__ == '__main__':
    main()