Spaces:
Sleeping
Sleeping
File size: 888 Bytes
1924502 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
import argparse, pathlib, torch, json
from agent import io, models, geometry
def main():
ap = argparse.ArgumentParser(description="Estimate object size from one image")
ap.add_argument("image", type=pathlib.Path)
ap.add_argument("--fx", type=float, required=True, help="focal length in pixels (x)")
ap.add_argument("--fy", type=float, required=True, help="focal length in pixels (y)")
args = ap.parse_args()
device = torch.device("cpu")
depth_net = models.load_depth(device)
mask_gen = models.load_sam(device)
img = io.load_rgb(args.image)
depth = models.predict_depth(depth_net, img, device)
masks = models.generate_masks(mask_gen, img)
mask = geometry.select_mask(masks)
stats = geometry.pixel_to_metric(mask["segmentation"], depth, args.fx, args.fy)
print(json.dumps(stats, indent=2))
if __name__ == "__main__":
main()
|