This repository has been archived by the owner on May 21, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset_models.py
104 lines (74 loc) · 3.1 KB
/
dataset_models.py
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import os
from typing import Any, Callable, Dict, List, Optional, Tuple
import pandas as pd
from PIL import Image
from torch import Tensor
from torch.utils.data import Dataset
from annotation_parser import BndBox, parse_xml
class ClassificationDataset(Dataset[Tensor]):
"""
Creates a dataset of cropped object images for classification.
"""
def __init__(self, anno_dir: str,
image_dir: str,
class_to_idx: Dict[str, int],
transform: Optional[Callable[..., Tensor]]) -> None:
self.image_dir = image_dir
self.datas: List[Tuple[str, str, BndBox]] = []
self.transform = transform
self.class_to_idx = class_to_idx
for image_file in os.listdir(image_dir):
anno_file = image_file.replace('jpg', 'xml')
anno_path = os.path.join(anno_dir, anno_file)
annotation = parse_xml(anno_path)
for obj in annotation['objects']:
file_name = image_file
name = obj['name']
bndbox = obj['bndbox']
self.datas.append((file_name, name, bndbox))
def __len__(self) -> int:
return len(self.datas)
def __getitem__(self, index: int) -> Tuple[Tensor, int]:
image_path = os.path.join(self.image_dir, self.datas[index][0])
with open(image_path, 'rb') as f:
image = Image.open(f)
image = image.convert('RGB')
bndbox = self.datas[index][2]
image = image.crop((bndbox['xmin'], bndbox['ymin'],
bndbox['xmax'], bndbox['ymax']))
if self.transform:
image = self.transform(image)
label = self.class_to_idx[self.datas[index][1]]
return image, label
class DetectionDataset(Dataset[Tensor]):
"""
Creates a dataset to evaluate the proposed bounding boxes.
"""
def __init__(self, image_dir: str,
prop_dir: str,
transform: Optional[Callable[..., Tensor]] = None) -> None:
self.transform = transform
self.item_list: List[Tuple[str, Any]] = []
for image_file in os.listdir(image_dir):
image_path = os.path.join(image_dir, image_file)
prop_path = os.path.join(
prop_dir, image_file.replace('jpg', 'csv'))
frame = pd.read_csv(prop_path)
# Crop images according to the proposals
for i in range(len(frame)):
prop = frame.iloc[i]
self.item_list.append((image_path, prop))
def __len__(self) -> int:
return len(self.item_list)
def __getitem__(self, index: int) -> Tuple[Tensor, int]:
image_path = self.item_list[index][0]
prop = self.item_list[index][1]
with open(image_path, 'rb') as f:
image = Image.open(f)
image = image.convert('RGB')
image = image.crop((prop['xmin'], prop['ymin'],
prop['xmax'], prop['ymax']))
label = int(prop['label'])
if self.transform:
image = self.transform(image)
return image, label