def test_predict(tmpdir, head):
model = ObjectDetector(num_classes=2, head=head, pretrained=False)
ds = DummyDetectionDataset((128, 128, 3), 1, 2, 10)
input_transform = IceVisionInputTransform()
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
dl = model.process_train_dataset(
ds,
2,
num_workers=0,
pin_memory=False,
input_transform=input_transform,
)
trainer.fit(model, dl)
dl = model.process_predict_dataset(
ds,
2,
input_transform=input_transform,
)
predictions = trainer.predict(model, dl, output="preds")
assert len(predictions[0][0]["bboxes"]) > 0
model.predict_kwargs = {"detection_threshold": 2}
predictions = trainer.predict(model, dl, output="preds")
assert len(predictions[0][0]["bboxes"]) == 0
def test_training(tmpdir, model):
model = ObjectDetector(num_classes=2,
model=model,
pretrained=False,
pretrained_backbone=False)
ds = DummyDetectionDataset((3, 224, 224), 1, 2, 10)
dl = DataLoader(ds, collate_fn=collate_fn)
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.fit(model, dl)
def test_init():
model = ObjectDetector(num_classes=2)
model.eval()
batch_size = 2
ds = DummyDetectionDataset((3, 224, 224), 1, 2, 10)
dl = DataLoader(ds, collate_fn=collate_fn, batch_size=batch_size)
data = next(iter(dl))
img = data[DefaultDataKeys.INPUT]
out = model(img)
assert len(out) == batch_size
assert {"boxes", "labels", "scores"} <= out[0].keys()
def test_init():
model = ObjectDetector(num_classes=2)
model.eval()
batch_size = 2
ds = DummyDetectionDataset((128, 128, 3), 1, 2, 10)
dl = model.process_predict_dataset(ds, batch_size=batch_size)
data = next(iter(dl))
out = model.forward(data[DefaultDataKeys.INPUT])
assert len(out) == batch_size
assert all(isinstance(res, dict) for res in out)
assert all("bboxes" in res for res in out)
assert all("labels" in res for res in out)
assert all("scores" in res for res in out)
def test_jit(tmpdir):
path = os.path.join(tmpdir, "test.pt")
model = ObjectDetector(2)
model.eval()
model = torch.jit.script(
model) # torch.jit.trace doesn't work with torchvision RCNN
torch.jit.save(model, path)
model = torch.jit.load(path)
out = model([torch.rand(3, 32, 32)])
# torchvision RCNN always returns a (Losses, Detections) tuple in scripting
out = out[1]
assert {"boxes", "labels", "scores"} <= out[0].keys()
def test_detection_fiftyone(tmpdir, head, backbone):
train_dataset = _create_synth_fiftyone_dataset(tmpdir)
datamodule = ObjectDetectionData.from_fiftyone(train_dataset=train_dataset, batch_size=1)
model = ObjectDetector(head=head, backbone=backbone, num_classes=datamodule.num_classes)
trainer = flash.Trainer(fast_dev_run=True, gpus=torch.cuda.device_count())
trainer.finetune(model, datamodule, strategy="freeze")
test_image_one = os.fspath(tmpdir / "test_one.png")
test_image_two = os.fspath(tmpdir / "test_two.png")
Image.new("RGB", (512, 512)).save(test_image_one)
Image.new("RGB", (512, 512)).save(test_image_two)
datamodule = ObjectDetectionData.from_files(predict_files=[str(test_image_one), str(test_image_two)], batch_size=1)
trainer.predict(model, datamodule=datamodule)
def test_detection_fiftyone(tmpdir, model, backbone):
train_dataset = _create_synth_fiftyone_dataset(tmpdir)
data = ObjectDetectionData.from_fiftyone(train_dataset=train_dataset,
batch_size=1)
model = ObjectDetector(model=model,
backbone=backbone,
num_classes=data.num_classes)
trainer = flash.Trainer(fast_dev_run=True)
trainer.finetune(model, data)
test_image_one = os.fspath(tmpdir / "test_one.png")
test_image_two = os.fspath(tmpdir / "test_two.png")
Image.new('RGB', (512, 512)).save(test_image_one)
Image.new('RGB', (512, 512)).save(test_image_two)
test_images = [str(test_image_one), str(test_image_two)]
model.predict(test_images)
def test_detection(tmpdir, head, backbone):
train_folder, coco_ann_path = _create_synth_coco_dataset(tmpdir)
data = ObjectDetectionData.from_coco(train_folder=train_folder,
train_ann_file=coco_ann_path,
batch_size=1)
model = ObjectDetector(head=head,
backbone=backbone,
num_classes=data.num_classes)
trainer = flash.Trainer(fast_dev_run=True, gpus=torch.cuda.device_count())
trainer.finetune(model, data, strategy="freeze")
test_image_one = os.fspath(tmpdir / "test_one.png")
test_image_two = os.fspath(tmpdir / "test_two.png")
Image.new("RGB", (512, 512)).save(test_image_one)
Image.new("RGB", (512, 512)).save(test_image_two)
test_images = [str(test_image_one), str(test_image_two)]
model.predict(test_images)
os.getcwd(), "data/coco128/annotations/instances_train2017.json"))
parser.add_argument('--max_epochs', type=int, default=1)
parser.add_argument('--learning_rate', type=float, default=1e-3)
parser.add_argument('--gpus', type=int, default=None)
args = parser.parse_args()
# 1. Download the data
if args.download:
# Dataset Credit: https://www.kaggle.com/ultralytics/coco128
download_data(
"https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip",
os.path.join(os.getcwd(), "data/"))
# 2. Load the Data
datamodule = ObjectDetectionData.from_coco(
train_folder=args.train_folder,
train_ann_file=args.train_ann_file,
batch_size=2)
# 3. Build the model
model = ObjectDetector(num_classes=datamodule.num_classes)
# 4. Create the trainer
trainer = flash.Trainer(max_epochs=args.max_epochs, gpus=args.gpus)
# 5. Finetune the model
trainer.finetune(model, datamodule)
# 6. Save it!
trainer.save_checkpoint("object_detection_model.pt")
import icedata # noqa: E402
# 1. Create the DataModule
data_dir = icedata.fridge.load_data()
datamodule = ObjectDetectionData.from_folders(
train_folder=data_dir,
predict_folder=data_dir,
val_split=0.1,
image_size=128,
parser=icedata.fridge.parser,
)
# 2. Build the task
model = ObjectDetector(head="efficientdet",
backbone="d0",
num_classes=datamodule.num_classes,
image_size=128)
# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=1)
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
# 4. Set the serializer and get some predictions
model.serializer = FiftyOneDetectionLabels(
return_filepath=True) # output FiftyOne format
predictions = trainer.predict(model, datamodule=datamodule)
predictions = list(chain.from_iterable(predictions)) # flatten batches
# 5. Visualize predictions in FiftyOne app
# Optional: pass `wait=True` to block execution until App is closed
session = visualize(predictions, wait=True)
def test_training(tmpdir, head):
model = ObjectDetector(num_classes=2, head=head, pretrained=False)
ds = DummyDetectionDataset((128, 128, 3), 1, 2, 10)
dl = model.process_train_dataset(ds, 2, 0, False, None)
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.fit(model, dl)
datamodule = ObjectDetectionData.from_icedata(
train_folder=data_dir,
predict_folder=os.path.join(data_dir, "odFridgeObjects", "images"),
val_split=0.1,
transform_kwargs={"image_size": 128},
parser=icedata.fridge.parser,
batch_size=8,
)
# 2. Build the task
model = ObjectDetector(
head="efficientdet",
backbone="d0",
labels=datamodule.labels,
image_size=128,
lr_scheduler=("multisteplr", {
"milestones": [20]
}),
)
# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=30)
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
# 4. Set the output and get some predictions
predictions = trainer.predict(model, datamodule=datamodule,
output="fiftyone") # output FiftyOne format
# 5. Visualize predictions in FiftyOne app
# Optional: pass `wait=True` to block execution until App is closed
from flash.image import ObjectDetectionData, ObjectDetector
# 1. Download the data
# Dataset Credit: https://www.kaggle.com/ultralytics/coco128
download_data(
"https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip",
"data/")
# 2. Load the Data
datamodule = ObjectDetectionData.from_coco(
train_folder="data/coco128/images/train2017/",
train_ann_file="data/coco128/annotations/instances_train2017.json",
val_split=0.3,
batch_size=4,
num_workers=4,
)
# 3. Build the model
model = ObjectDetector(model="retinanet", num_classes=datamodule.num_classes)
# 4. Create the trainer
trainer = flash.Trainer(max_epochs=3,
limit_train_batches=1,
limit_val_batches=1)
# 5. Finetune the model
trainer.finetune(model, datamodule=datamodule)
# 6. Save it!
trainer.save_checkpoint("object_detection_model.pt")