first commit

ldm/modules/midas/midas/base_model.py

+import torch
+
+
+class BaseModel(torch.nn.Module):
+    def load(self, path):
+        """Load model from file.
+
+        Args:
+            path (str): file path
+        """
+        parameters = torch.load(path, map_location=torch.device('cpu'))
+
+        if "optimizer" in parameters:
+            parameters = parameters["model"]
+
+        self.load_state_dict(parameters)

model/mixins.py

+from typing import overload, Any, Dict
+import torch
+
+
+class ImageLoggerMixin:
+
+    @overload
+    def log_images(self, batch: Any, **kwargs: Dict[str, Any]) -> Dict[str, torch.Tensor]:
+        ...

scripts/make_file_list.py

+import sys
+sys.path.append(".")
+import os
+from argparse import ArgumentParser
+
+from utils.file import list_image_files
+
+
+parser = ArgumentParser()
+parser.add_argument("--img_folder", type=str, required=True)
+parser.add_argument("--val_size", type=int, required=True)
+parser.add_argument("--save_folder", type=str, required=True)
+parser.add_argument("--follow_links", action="store_true")
+args = parser.parse_args()
+
+files = list_image_files(
+    args.img_folder, exts=(".jpg", ".png", ".jpeg"), follow_links=args.follow_links,
+    log_progress=True, log_every_n_files=10000
+)
+
+print(f"find {len(files)} images in {args.img_folder}")
+assert args.val_size < len(files)
+
+val_files = files[:args.val_size]
+train_files = files[args.val_size:]
+
+os.makedirs(args.save_folder, exist_ok=True)
+
+with open(os.path.join(args.save_folder, "train.list"), "w") as fp:
+    for file_path in train_files:
+        fp.write(f"{file_path}\n")
+
+with open(os.path.join(args.save_folder, "val.list"), "w") as fp:
+    for file_path in val_files:
+        fp.write(f"{file_path}\n")