Fix JSON prediction coordinate exports

Normalize all exported prediction JSON to original input-image coordinates for both legacy TF and PyTorch pipelines, including image size metadata and nulls for hidden keypoints. This also threads the source and annotated image sizes through the app, updates the Gradio SDK version, and removes the extra README demo examples.

README.md

@@ -5,7 +5,7 @@ colorFrom: blue
 colorTo: purple
 sdk: gradio
 python_version: 3.12
-sdk_version: 6.11.0
+sdk_version: 6.29.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -106,7 +106,13 @@ def predict_pipeline_pytorch(
 
     pose_model, detector = PYTORCH_MODELS[superanimal]
     download_file = save_results_pytorch(
-        animals, map_label_id_to_str, superanimal, pose_model, None if flag_dlc_only else detector
+        animals,
+        map_label_id_to_str,
+        superanimal,
+        pose_model,
+        None if flag_dlc_only else detector,
+        image_size=img_input.size,
+        annotated_size=img_output.size,
     )
     return finalize_outputs(
         img_output,
@@ -220,7 +226,9 @@ def predict_pipeline(
             colormap=colormap,
         )
 
-        donw_file = save_results_only_dlc(list_kpts_per_crop[0], map_label_id_to_str, dlc_model_name)
+        donw_file = save_results_only_dlc(
+            list_kpts_per_crop[0], map_label_id_to_str, dlc_model_name, image_size=img_input.size
+        )
 
         return finalize_outputs(
             img_input, donw_file, [list_kpts_per_crop[0]], map_label_id_to_str, flag_color_by_confidence, colormap
@@ -260,7 +268,13 @@ def predict_pipeline(
 
         # Save detection results as json
         download_file = save_results_as_json(
-            md_results, list_kpts_per_crop, list_bboxes, map_label_id_to_str, dlc_model_name, mega_model_input
+            md_results,
+            list_kpts_per_crop,
+            list_bboxes,
+            map_label_id_to_str,
+            dlc_model_name,
+            mega_model_input,
+            image_size=img_input.size,
         )
 
         return finalize_outputs(

ui_utils.py

@@ -221,8 +221,6 @@ def gradio_description_and_examples():
         for image in (
             "examples/dog.jpeg",
             "examples/cat.jpg",
-            "examples/lynx.jpg",
-            "examples/giraffe.jpg",
         )
         for font_size, marker_size in [example_sizes(image)]
     ]

viz_utils.py

@@ -1,4 +1,5 @@
 import json
+import math
 from datetime import date
 
 import numpy as np
@@ -161,6 +162,36 @@ def label_text_color(background):
 
 
 ###########################################
+#  JSON outputs: pixel coordinates in the input image, hidden keypoints as null
+COORDINATES = "pixels in the input image (after EXIF orientation), origin top-left, x right, y down"
+
+
+def keypoints_to_json(kpts, offset=(0, 0), scale=(1.0, 1.0)):
+    """[x, y, confidence] per keypoint, mapped by (k + offset) * scale; NaN (below threshold) becomes null."""
+    out = []
+    for x, y, conf in np.asarray(kpts, dtype=float)[:, :3]:
+        if np.isnan(conf):
+            out.append(None)
+        else:
+            out.append([(x + offset[0]) * scale[0], (y + offset[1]) * scale[1], conf])
+    return out
+
+
+def write_json(info, path_to_output_file):
+    with open(path_to_output_file, "w") as f:
+        json.dump(info, f, indent=1, allow_nan=False)
+    return path_to_output_file
+
+
+def json_header(image_size, annotated_size):
+    return {
+        "date": str(today),
+        "coordinates": COORDINATES,
+        "image_size": list(image_size),
+        "annotated_image_size": list(annotated_size),
+    }
+
+
 def save_results_as_json(
     md_results,
     dlc_outputs,
@@ -168,110 +199,79 @@ def save_results_as_json(
     map_dlc_label_id_to_str,
     model,
     mega_model_input,
+    image_size,
     path_to_output_file="download_predictions.json",
 ):
-    """
-    Output detections as json file
+    """TF (legacy) MegaDetector + DLC results.
 
-    animal_bboxes: detection rows [x1,y1,x2,y2,conf,label], one per entry of dlc_outputs
+    animal_bboxes: detection rows [x1,y1,x2,y2,conf,label] in the MegaDetector frame, one per entry of dlc_outputs
+    dlc_outputs: keypoints relative to each crop (crops start at floor(x1), floor(y1), see crop_animal_detections)
     """
-    # initialise dict to save to json
-    info = {}
-    info["date"] = str(today)
+    md_h, md_w = md_results.ims[0].shape[:2]
+    scale = (image_size[0] / md_w, image_size[1] / md_h)
+    info = json_header(image_size, (md_w, md_h))
     info["MD_model"] = str(mega_model_input)
-    # info from megaDetector
-    info["file"] = md_results.files[0]
-    number_bb = len(md_results.xyxy[0].tolist())
-    info["number_of_bb"] = number_bb
-    # info from DLC
+    info["number_of_bb"] = len(dlc_outputs)
     info["dlc_model"] = model
-    labels = [n for n in map_dlc_label_id_to_str.values()]
-
-    # define aux dict for every animal bounding box above threshold
-    for i in range(len(dlc_outputs)):
-        aux = {}
-        # MD output
-        corner_x1, corner_y1, corner_x2, corner_y2, confidence, _ = animal_bboxes[i]
-        aux["corner_1"] = (corner_x1, corner_y1)
-        aux["corner_2"] = (corner_x2, corner_y2)
-        aux["predict MD"] = md_results.names[0]
-        aux["confidence MD"] = confidence
-
-        # DLC output
-        kypts = []
-        for s in dlc_outputs[i]:
-            aux1 = []
-            for j in s:
-                aux1.append(float(j))
-
-            kypts.append(aux1)
-        aux["dlc_pred"] = dict(zip(labels, kypts, strict=True))
-        info["bb_" + str(i)] = aux
-
-    # save dict as json
-    with open(path_to_output_file, "w") as f:
-        json.dump(info, f, indent=1)
-        print(f"Output file saved at {path_to_output_file}")
-
-    return path_to_output_file
-
-
-def save_results_only_dlc(dlc_outputs, map_label_id_to_str, model, output_file="dowload_predictions_dlc.json"):
-    """
-    write json dlc output
-    """
-    info = {}
-    info["date"] = str(today)
-    labels = [n for n in map_label_id_to_str.values()]
+    labels = list(map_dlc_label_id_to_str.values())
+
+    for i, kpts in enumerate(dlc_outputs):
+        x1, y1, x2, y2, confidence, _ = animal_bboxes[i]
+        info["bb_" + str(i)] = {
+            "corner_1": (x1 * scale[0], y1 * scale[1]),
+            "corner_2": (x2 * scale[0], y2 * scale[1]),
+            "predict MD": md_results.names[0],
+            "confidence MD": float(confidence),
+            "dlc_pred": dict(
+                zip(labels, keypoints_to_json(kpts, offset=(math.floor(x1), math.floor(y1)), scale=scale), strict=True)
+            ),
+        }
+    return write_json(info, path_to_output_file)
+
+
+def save_results_only_dlc(
+    dlc_outputs, map_label_id_to_str, model, image_size, output_file="dowload_predictions_dlc.json"
+):
+    """TF (legacy) DLC run on the whole input image (keypoints already in input pixels)."""
+    info = json_header(image_size, image_size)
     info["dlc_model"] = model
-    kypts = []
-    for s in dlc_outputs:
-        aux1 = []
-        for j in s:
-            aux1.append(float(j))
-
-        kypts.append(aux1)
-    info["dlc_pred"] = dict(zip(labels, kypts, strict=True))
-
-    with open(output_file, "w") as f:
-        json.dump(info, f, indent=1)
-        print(f"Output file saved at {output_file}")
-
-    return output_file
+    info["dlc_pred"] = dict(zip(map_label_id_to_str.values(), keypoints_to_json(dlc_outputs), strict=True))
+    return write_json(info, output_file)
 
 
 def save_results_pytorch(
-    animals, map_label_id_to_str, model, pose_model, detector, path_to_output_file="download_predictions.json"
+    animals,
+    map_label_id_to_str,
+    model,
+    pose_model,
+    detector,
+    image_size,
+    annotated_size,
+    path_to_output_file="download_predictions.json",
 ):
-    """
-    Output PyTorch SuperAnimal predictions as json file (same layout as save_results_as_json)
+    """PyTorch SuperAnimal results (same layout as save_results_as_json).
 
-    animals: list of {'bbox': [x1,y1,x2,y2,conf], 'kpts': (num_keypoints, 3)}, in image coords
+    animals: list of {'bbox': [x1,y1,x2,y2,conf], 'kpts': (num_keypoints, 3)}, in the annotated (resized) image
     detector: None if the detector was skipped (whole image used as one animal)
     """
-    info = {}
-    info["date"] = str(today)
+    scale = (image_size[0] / annotated_size[0], image_size[1] / annotated_size[1])
+    info = json_header(image_size, annotated_size)
     info["backend"] = "pytorch"
     info["dlc_model"] = model
     info["pose_model"] = pose_model
     info["detector"] = detector
     info["number_of_bb"] = len(animals)
-    labels = [n for n in map_label_id_to_str.values()]
+    labels = list(map_label_id_to_str.values())
 
     for i, animal in enumerate(animals):
-        corner_x1, corner_y1, corner_x2, corner_y2, confidence = animal["bbox"]
-        aux = {}
-        aux["corner_1"] = (corner_x1, corner_y1)
-        aux["corner_2"] = (corner_x2, corner_y2)
-        aux["confidence"] = confidence
-        aux["dlc_pred"] = dict(zip(labels, [[float(v) for v in kpt] for kpt in animal["kpts"]], strict=True))
-        info["bb_" + str(i)] = aux
-
-    with open(path_to_output_file, "w") as f:
-        json.dump(info, f, indent=1)
-        print(f"Output file saved at {path_to_output_file}")
-
-    return path_to_output_file
+        x1, y1, x2, y2, confidence = animal["bbox"]
+        info["bb_" + str(i)] = {
+            "corner_1": (x1 * scale[0], y1 * scale[1]),
+            "corner_2": (x2 * scale[0], y2 * scale[1]),
+            "confidence": confidence,
+            "dlc_pred": dict(zip(labels, keypoints_to_json(animal["kpts"], scale=scale), strict=True)),
+        }
+    return write_json(info, path_to_output_file)
 
 
 ###########################################