|
|
| import os |
| import torch |
| from torch.utils.data import Dataset |
| from PIL import Image |
|
|
| class WiderFaceMultiBoxDataset(Dataset): |
| def __init__(self, images_dir, annotation_file, S=20, B=3, transform=None): |
| self.images_dir = images_dir |
| self.annotation_file = annotation_file |
| self.transform = transform |
| self.S = S |
| self.B = B |
| self.data = self.parse_annotations() |
|
|
| def parse_annotations(self): |
| with open(self.annotation_file, "r") as f: |
| lines = [l.strip() for l in f] |
| data = [] |
| i = 0 |
| while i < len(lines): |
| fn = lines[i]; i += 1 |
| num = int(lines[i]); i += 1 |
| boxes = [] |
| for _ in range(num): |
| x,y,w,h = map(int, lines[i].split()[:4]); i += 1 |
| if w>0 and h>0: boxes.append((x,y,w,h)) |
| if boxes: data.append((fn, boxes)) |
| return data |
|
|
| def encode_target(self, boxes, W, H): |
| t = torch.zeros((self.S,self.S,self.B*5)) |
| for x,y,w,h in boxes: |
| cx=(x+w/2)/W; cy=(y+h/2)/H; wn=w/W; hn=h/H |
| gx=int(cx*self.S); gy=int(cy*self.S) |
| if 0<=gx<self.S and 0<=gy<self.S: |
| for b in range(self.B): |
| if t[gy,gx,b*5+4]==0: |
| t[gy,gx,b*5:b*5+5]=torch.tensor([cx,cy,wn,hn,1.0]) |
| break |
| return t |
|
|
| def __getitem__(self, idx): |
| fn, boxes = self.data[idx] |
| img = Image.open(os.path.join(self.images_dir, fn)).convert("RGB") |
| W,H = img.size |
| if self.transform: img = self.transform(img) |
| return img, self.encode_target(boxes,W,H) |
|
|
| def __len__(self): return len(self.data) |
|
|
