How can I create a useful labelling function for images from a directory?

Question

I am working on Jupyter Notebook with the fastai DataBlock and Dataloaders API to prepare batches for a neural network. Currently what I've done is:

'''

path= Path('TensorFlow\workspace\images\mini_alphabet')''' 

#created this path to the folder containing 20 images labelled A1,A2,A3,etc and 20 labelled B,B2,B3 etc

and when i use.ls() I can open and view the output in this form:

[Path('A1.jpg'),Path('A10.jpg'),Path('A11.jpg'),Path('A12.jpg'),Path('A13.jpg'),Path('A14.jpg'),Path('A15.jpg'),Path('A16.jpg'),Path('A17.jpg'),Path('A18.jpg')...]

What I want to do is make a labelling function that iterates through the mini_alphabet folder, checks to see if the image name starts with the letter A or B, and then returns that letter as the label. So far I've written this function:

def label_alphabet(fname):
    labels = os.listdir(fname)  # <this outputs a list ['A1.jpg','A2.jpg',...]
    for l in labels:
        if l[0].startswith('A'):
            return "A"
        else:
            return "B"

Unfortunately when I use this it seems to label every single image as A. What should I do differently here? Also if I wanted to apply a larger number of labels (ranging from AG), how exactly should I set up that code so it labels through all of them. I was thinking of iterating '''for i''' through a list of letters and returning labels that match the filename.

Thanks for the help!

Answer 1

try if l.split(".")[0].startwith("A") instead of if l[0].startswith('A'):

If you want an array of letters you could use list(string.ascii_uppercase)

Also put a tab between for and the code that follows from there