Spaces:
Sleeping
Sleeping
File size: 1,599 Bytes
6e89f30 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import glob
import cv2
import pandas as pd
import torch
import os
from src.config import cfg
from dataclasses import dataclass
@dataclass
class CaptchaDataset(torch.utils.data.Dataset):
def __init__(self,folder:str):
self.data_root = cfg.data_root
df = pd.read_csv(f"{self.data_root}/{folder}/labels.csv")
self.data = []
for _,row in df.iterrows():
filename = row['filename']
label = row['label']
img_path = f"{self.data_root}/{folder}/{row['filename']}"
# Check if file actually exists
if os.path.exists(img_path):
self.data.append((img_path,label,folder))
else:
print(f"Warning: Image file not found: {img_path}")
print(f"Loaded {len(self.data)} valid images from {folder}")
self.img_dim = (cfg.W_max, cfg.H) # cv2.resize expects (width, height)
def __len__(self):
return len(self.data)
def __getitem__(self,idx):
img_path, label_string,folder = self.data[idx]
# Load image with error checking
img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE if cfg.grayscale else cv2.IMREAD_COLOR)
if img is None:
raise ValueError(f"Failed to load image: {img_path}")
img = cv2.resize(img, self.img_dim)
img_tensor = torch.from_numpy(img).float()/255.0 # Normalize to [0,1]
img_tensor = img_tensor.unsqueeze(0) # Add channel dimension
return img_tensor, label_string, img_path
|