[FIXED] How can I iterate over all the batches from DataLoader?

Issue

I wanted to iterate over all the batches and save the images but with this process its saving only images of the first batch

    for batch_idx, (test_data, test_targets) in enumerate(test_loader):
        for i in range(0, test_loader.batch_size-1):
            img = np.array(test_data[i][0])*255
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            low_black = np.array([0,0,0])
            high_black = np.array([360,255,0])
            mask = cv2.inRange(hsv, low_black, high_black)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img[mask>0]=random.choice(list(color_dict.values()))
            cv2.imwrite(f'/content/test_data/{test_targets[i].item()}_{i+1}.png', img)

Solution

Since i is start from 0 to batch_size at every batch so the saved names are duplicated. One common way to solve it is using count:

count = 0 # here
for batch_idx, (test_data, test_targets) in enumerate(test_loader):
    for i in range(0, test_loader.batch_size-1):
        img = np.array(test_data[i][0])*255
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        low_black = np.array([0,0,0])
        high_black = np.array([360,255,0])
        mask = cv2.inRange(hsv, low_black, high_black)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img[mask>0]=random.choice(list(color_dict.values()))
        cv2.imwrite(f'/content/test_data/{test_targets[i].item()}_{count}.png', img)
        count+=1 # plus one every sample

Answered By - CuCaRot