Spaces:
Running
Running
File size: 3,161 Bytes
09801ca | 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 | import os
import zipfile
import json
import numpy as np
import cv2
from pathlib import Path
out_dir = Path("sample_cv_test_datasets")
out_dir.mkdir(exist_ok=True)
def create_sample_img(filename, color=(0, 200, 100)):
img = np.zeros((300, 300, 3), dtype=np.uint8)
cv2.rectangle(img, (50, 50), (250, 250), color, -1)
cv2.putText(img, "TEST DATASET", (60, 160), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
cv2.imwrite(filename, img)
# 1. Instance Segmentation ZIP Dataset
seg_dir = out_dir / "sample_instance_segmentation"
seg_dir.mkdir(exist_ok=True)
create_sample_img(str(seg_dir / "car_sample.jpg"), (220, 50, 50))
coco_seg = {
"images": [{"id": 1, "file_name": "car_sample.jpg", "height": 300, "width": 300}],
"categories": [{"id": 1, "name": "vehicle_segmentation"}],
"annotations": [{
"id": 1, "image_id": 1, "category_id": 1,
"segmentation": [[50, 50, 250, 50, 250, 250, 50, 250]],
"bbox": [50, 50, 200, 200], "area": 40000, "iscrowd": 0
}]
}
with open(seg_dir / "_annotations.coco.json", "w") as f:
json.dump(coco_seg, f, indent=2)
with zipfile.ZipFile(out_dir / "Instance_Segmentation_Test_Dataset.zip", "w") as z:
z.write(seg_dir / "car_sample.jpg", "car_sample.jpg")
z.write(seg_dir / "_annotations.coco.json", "_annotations.coco.json")
# 2. Keypoint Detection ZIP Dataset
pose_dir = out_dir / "sample_keypoint_pose"
pose_dir.mkdir(exist_ok=True)
create_sample_img(str(pose_dir / "person_pose.jpg"), (50, 100, 250))
coco_pose = {
"images": [{"id": 1, "file_name": "person_pose.jpg", "height": 300, "width": 300}],
"categories": [{"id": 1, "name": "human_pose", "keypoints": ["nose", "left_eye", "right_eye", "left_shoulder", "right_shoulder"]}],
"annotations": [{
"id": 1, "image_id": 1, "category_id": 1,
"keypoints": [150, 100, 2, 140, 90, 2, 160, 90, 2, 100, 180, 2, 200, 180, 2],
"bbox": [50, 50, 200, 200], "num_keypoints": 5, "area": 40000, "iscrowd": 0
}]
}
with open(pose_dir / "_annotations.coco.json", "w") as f:
json.dump(coco_pose, f, indent=2)
with zipfile.ZipFile(out_dir / "Keypoint_Detection_Test_Dataset.zip", "w") as z:
z.write(pose_dir / "person_pose.jpg", "person_pose.jpg")
z.write(pose_dir / "_annotations.coco.json", "_annotations.coco.json")
# 3. OCR Text Extraction ZIP Dataset
ocr_dir = out_dir / "sample_ocr_dataset"
ocr_dir.mkdir(exist_ok=True)
create_sample_img(str(ocr_dir / "ocr_text.jpg"), (10, 180, 220))
coco_ocr = {
"images": [{"id": 1, "file_name": "ocr_text.jpg", "height": 300, "width": 300}],
"categories": [{"id": 1, "name": "text_block"}],
"annotations": [{
"id": 1, "image_id": 1, "category_id": 1,
"text": "TEST DATASET",
"bbox": [60, 140, 180, 40], "area": 7200, "iscrowd": 0
}]
}
with open(ocr_dir / "_annotations.coco.json", "w") as f:
json.dump(coco_ocr, f, indent=2)
with zipfile.ZipFile(out_dir / "OCR_Text_Test_Dataset.zip", "w") as z:
z.write(ocr_dir / "ocr_text.jpg", "ocr_text.jpg")
z.write(ocr_dir / "_annotations.coco.json", "_annotations.coco.json")
print("Generated sample benchmark test datasets cleanly.")
|