Spaces:
Runtime error
Runtime error
File size: 4,329 Bytes
0a216c0 | 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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """
Run new models (emanuskript, catmus, zone) and convert predictions to COCO format.
Uses the same logic as app.py and test_combined_models.py.
"""
import os
import sys
import json
from pathlib import Path
import tempfile
# Add project root to path
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(os.path.dirname(SCRIPT_DIR))
sys.path.insert(0, PROJECT_ROOT)
# Import from project root
from test_combined_models import run_model_predictions, combine_and_filter_predictions
def run_new_models_on_image(image_path, conf_threshold=0.25, iou_threshold=0.45):
"""
Run new models on a single image and return COCO format predictions.
Args:
image_path: Path to image file
conf_threshold: Confidence threshold (not directly used, but kept for consistency)
iou_threshold: IoU threshold (not directly used, but kept for consistency)
Returns:
COCO format dictionary with predictions
"""
with tempfile.TemporaryDirectory() as tmp_dir:
# Run 3 models and save predictions JSON
labels_folders = run_model_predictions(image_path, tmp_dir)
# Combine & filter to coco_class_mapping
coco_json = combine_and_filter_predictions(
image_path, labels_folders, output_json_path=None
)
return coco_json
def process_dataset(images_dir, output_dir, conf_threshold=0.25, iou_threshold=0.45):
"""
Process all images in a directory with new models.
Args:
images_dir: Directory containing images
output_dir: Directory to save COCO JSON files
conf_threshold: Confidence threshold
iou_threshold: IoU threshold
Returns:
Merged COCO format dictionary for all images
"""
# Get all image files
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'}
image_files = [
f for f in os.listdir(images_dir)
if os.path.splitext(f)[1].lower() in image_extensions
]
all_coco_dicts = []
image_id = 1
for image_file in sorted(image_files):
image_path = os.path.join(images_dir, image_file)
print(f"Processing {image_file}...")
try:
coco = run_new_models_on_image(
image_path,
conf_threshold=conf_threshold,
iou_threshold=iou_threshold
)
# Update image ID if needed
if len(coco["images"]) > 0:
coco["images"][0]["id"] = image_id
# Update annotation image_ids
for ann in coco["annotations"]:
ann["image_id"] = image_id
all_coco_dicts.append(coco)
image_id += 1
# Save individual file
output_path = os.path.join(output_dir, f"{Path(image_file).stem}_new.json")
with open(output_path, 'w') as f:
json.dump(coco, f, indent=2)
except Exception as e:
print(f"Error processing {image_file}: {e}")
import traceback
traceback.print_exc()
continue
# Merge all COCO dicts
merged = {
"info": {"description": "New models predictions - merged"},
"licenses": [],
"images": [],
"annotations": [],
"categories": []
}
# Use categories from first COCO (they should all be the same)
if len(all_coco_dicts) > 0:
merged["categories"] = all_coco_dicts[0]["categories"]
# Merge images and annotations
ann_id = 1
for coco in all_coco_dicts:
merged["images"].extend(coco["images"])
for ann in coco["annotations"]:
new_ann = ann.copy()
new_ann["id"] = ann_id
merged["annotations"].append(new_ann)
ann_id += 1
return merged
if __name__ == "__main__":
# Test on single image
test_image = "../../e-codices_bbb-0219_044r_max.jpg"
coco = run_new_models_on_image(test_image)
print(f"Predictions: {len(coco['annotations'])} annotations")
print(f"Categories: {[c['name'] for c in coco['categories']]}")
with open("test_new_models.json", "w") as f:
json.dump(coco, f, indent=2)
|