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()