Spaces:
Build error
Build error
| import pytorch_lightning as pl | |
| from det_models.model import POIDetection | |
| from datasets_signboard_detection.datamodule import POIDataModule | |
| from det_models.inference_signboard_detection import POIDetectionTask | |
| def load_model(checkpoint_path): | |
| model = POIDetection.load_from_checkpoint(checkpoint_path=checkpoint_path) | |
| return model | |
| def inference_signboard(image, checkpoint): | |
| dm = POIDataModule(data=image, seed=42) | |
| dm.setup("predict") | |
| model = load_model(checkpoint) | |
| task = POIDetectionTask(model) | |
| # accelerator='gpu', devices=1 | |
| trainer = pl.Trainer(gpus=0, max_epochs=-1) | |
| trainer.predict(task, datamodule=dm) | |
| return task.output | |