def test_detection_fiftyone(tmpdir, head, backbone):
train_dataset = _create_synth_fiftyone_dataset(tmpdir)
data = ObjectDetectionData.from_fiftyone(train_dataset=train_dataset,
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)
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)
"https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip",
"data/")
datamodule = ObjectDetectionData.from_coco(
train_folder="data/coco128/images/train2017/",
train_ann_file="data/coco128/annotations/instances_train2017.json",
val_split=0.1,
image_size=128,
)
# 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. Detect objects in a few images!
predictions = model.predict([
"data/coco128/images/train2017/000000000625.jpg",
"data/coco128/images/train2017/000000000626.jpg",
"data/coco128/images/train2017/000000000629.jpg",
])
print(predictions)
# 5. Save the model!
trainer.save_checkpoint("object_detection_model.pt")