Zitti123's picture
πŸš€ Initial deployment: Smart Object Size Estimator with AI-powered depth estimation and SAM segmentation
1924502
raw
history blame contribute delete
888 Bytes
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()