Set up the prediction environment. This method initializes the model, its parameters, and the GPU device for computation. It also loads the STMFNet model using the specified checkpoint. Lastly, it ensures there is an output directory for storing the enhanced videos.
(self)
| 35 | |
| 36 | class Predictor(BasePredictor): |
| 37 | def setup(self): |
| 38 | """ |
| 39 | Set up the prediction environment. |
| 40 | |
| 41 | This method initializes the model, its parameters, and the GPU device for computation. |
| 42 | It also loads the STMFNet model using the specified checkpoint. |
| 43 | Lastly, it ensures there is an output directory for storing the enhanced videos. |
| 44 | """ |
| 45 | if not os.path.exists(STMFNET_WEIGHTS_PATH): |
| 46 | download_weights(STMFNET_WEIGHTS_URL, STMFNET_WEIGHTS_PATH) |
| 47 | |
| 48 | args = SimpleNamespace( |
| 49 | **{ |
| 50 | "gpu_id": (gpu_id := 0), |
| 51 | "net": (net := "STMFNet"), |
| 52 | "checkpoint": (checkpoint := STMFNET_WEIGHTS_PATH), |
| 53 | "size": (size := "1920x1080"), |
| 54 | "patch_size": (patch_size := None), |
| 55 | "overlap": (overlap := None), |
| 56 | "batch_size": (batch_size := None), |
| 57 | "out_fps": (out_fps := 144), |
| 58 | "out_dir": (out_dir := "."), |
| 59 | "featc": (featc := [64, 128, 256, 512]), |
| 60 | "featnet": (featnet := "UMultiScaleResNext"), |
| 61 | "featnorm": (featnorm := "batch"), |
| 62 | "kernel_size": (kernel_size := 5), |
| 63 | "dilation": (dilation := 1), |
| 64 | "finetune_pwc": (finetune_pwc := False), |
| 65 | } |
| 66 | ) |
| 67 | torch.cuda.set_device(gpu_id) |
| 68 | |
| 69 | self.net = net |
| 70 | self.size = size |
| 71 | self.model = STMFNet(args).cuda() |
| 72 | print("Loading the model...") |
| 73 | checkpoint = torch.load(checkpoint) |
| 74 | self.model.load_state_dict(checkpoint["state_dict"]) |
| 75 | self.model.eval() |
| 76 | |
| 77 | if not os.path.exists(out_dir): |
| 78 | os.makedirs(out_dir) |
| 79 | |
| 80 | def predict( |
| 81 | self, |
nothing calls this directly
no test coverage detected