AI_PROJECT/Image.py

import pygame
import displayControler as dCon
import random
import neuralnetwork
import os

class Image:
    def __init__(self):
        self.plants_image_dict={}
        self.tractor_image=None
        self.garage_image=None
        self.stone_image=None
        self.gasStation_image=None
    def load_images(self):
        files_plants={
            0:"borowka",
            1:"kukurydza",
            2:"pszenica",
            3:"slonecznik",
            4:"winogrono",
            5:"ziemniak",
            6:"dirt",
            7:"mud",
            8:"road"}
        for index in files_plants:
            if index >= 6:
                plant_image = pygame.image.load("images/" + files_plants[index] + ".jpg")
            else:
                plant_image=pygame.image.load("images/plants/"+files_plants[index]+".jpg")
            plant_image=pygame.transform.scale(plant_image,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))
            self.plants_image_dict[files_plants[index]]=plant_image
        tractor_image=pygame.image.load("images/traktor.png")
        tractor_image=pygame.transform.scale(tractor_image,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))
        garage=pygame.image.load("images/garage.png")
        self.garage_image=pygame.transform.scale(garage,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))
        stone=pygame.image.load("images/stone.png")
        self.stone_image=pygame.transform.scale(stone,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))
        gasStation=pygame.image.load("images/gasStation.png")
        self.gasStation_image=pygame.transform.scale(gasStation,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))

    def return_random_plant(self):
        x=random.randint(0,5) #disabled dirt and mud generation
        keys=list(self.plants_image_dict.keys())
        plant=keys[x]
        return (plant,self.plants_image_dict[plant])
    
    def return_plant(self,plant_name):
        return (plant_name,self.plants_image_dict[plant_name])
    
    def return_garage(self):
        return self.garage_image

    def return_stone(self):
        return self.stone_image
    
    def return_gasStation(self):
        return self.gasStation_image

# losowanie zdjęcia z testowego datasetu bez powtórzeń
imagePathList = []
def getRandomImageFromDataBase():
    label = random.choice(neuralnetwork.labels)
    folderPath = f"dataset/test/{label}"
    files = os.listdir(folderPath)
    random_image = random.choice(files)
    imgPath = os.path.join(folderPath, random_image)

    while imgPath in imagePathList:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                quit()
        label = random.choice(neuralnetwork.labels)
        folderPath = f"dataset/test/{label}"
        files = os.listdir(folderPath)
        random_image = random.choice(files)
        imgPath = os.path.join(folderPath, random_image)

    imagePathList.append(imgPath)

    image = pygame.image.load(imgPath)
    image=pygame.transform.scale(image,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))
    return image, label, imgPath

def getSpedifiedImageFromDatabase(label):
    folderPath = f"dataset/test/{label}"
    files = os.listdir(folderPath)
    random_image = random.choice(files)
    imgPath = os.path.join(folderPath, random_image)

    while imgPath in imagePathList:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                quit()
        label = random.choice(neuralnetwork.labels)
        folderPath = f"dataset/test/{label}"
        files = os.listdir(folderPath)
        random_image = random.choice(files)
        imgPath = os.path.join(folderPath, random_image)

    imagePathList.append(imgPath)

    image = pygame.image.load(imgPath)
    image=pygame.transform.scale(image,(dCon.CUBE_SIZE,dCon.CUBE_SIZE))
    return image, label, imgPath