github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

reuse existing prepare_target

Browse Source
This commit is contained in:
Francis Lata
2024-10-05 18:21:15 -07:00
parent 30f6f6a094
commit 8ca848d542
2 changed files with 2 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
examples/mlperf/dataloader.py
View File

@@ -357,12 +357,12 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
### RetinaNet
def load_retinanet_data(base_dir:Path, queue_in:Queue, queue_out:Queue, X:Tensor):
from extra.datasets.openimages import image_load, transform_coco_polys_to_mask
from extra.datasets.openimages import image_load, prepare_target
while (data:=queue_in.get()) is not None:
idx, img, ann = data
img_id = img["id"]
img, img_size = image_load(base_dir, "train", img["file_name"]) # TODO: resize this with the target!
img, tgt = transform_coco_polys_to_mask(img, img_id, img_size, ann)
tgt = prepare_target(ann, img_id, img_size)
X[idx].contiguous().realize().lazydata.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()

18
extra/datasets/openimages.py
View File

@@ -179,24 +179,6 @@ def download_dataset(base_dir:Path, subset:str) -> Path:
return ann_file
def transform_coco_polys_to_mask(img, img_id, img_size, ann, filter_is_crowd:bool=True):
w, h = img_size
if filter_is_crowd: ann = [obj for obj in ann if obj["iscrowd"] == 0]
boxes = np.array([obj["bbox"] for obj in ann])
boxes = boxes.reshape(-1, 4)
boxes[:, 2:] += boxes[:, :2]
boxes[:, 0::2].clip(0, w)
boxes[:, 1::2].clip(0, h)
classes = np.array([obj["category_id"] for obj in ann])
keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
boxes, classes = boxes[keep], classes[keep]
area, is_crowd = [obj["area"] for obj in ann], [obj["iscrowd"] for obj in ann]
target = {"boxes": boxes, "labels": classes, "image_id": img_id, "area": area, "iscrowd": is_crowd}
return img, target
if __name__ == "__main__":
download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), "train")