Spaces:
Sleeping
Sleeping
File size: 2,798 Bytes
fa83aac | 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 | import os
import json
import itertools
from datasets import load_dataset
def fetch_real_oct_images():
output_dir = "images"
os.makedirs(output_dir, exist_ok=True)
# Label mapping in the keremberke dataset
labels_map = {0: "CNV", 1: "DME", 2: "DRUSEN", 3: "NORMAL"}
print("Connecting to Hugging Face Cloud to stream real OCT images...)")
print("This will only download ~5MB instead of 12GB!")
try:
# Streaming=True ensures we don't download the zip!
dataset = load_dataset("keremberke/oct-image-classification", "full", split="train", streaming=True)
except Exception as e:
print(f"Error connecting to dataset: {e}")
return
ground_truth = {}
counts = {"CNV": 0, "DME": 0, "DRUSEN": 0, "NORMAL": 0}
target_per_class = 10 # 40 images total (10 per class)
for item in dataset:
label_id = item["labels"]
label_name = labels_map.get(label_id, "NORMAL")
if counts[label_name] < target_per_class:
img = item["image"]
filename = f"{label_name}_{counts[label_name] + 1}.jpg"
filepath = os.path.join(output_dir, filename)
# Save the image
img.save(filepath)
# Formulate the ground truth metadata
keywords = []
if label_name == "CNV": keywords = ["subretinal fluid", "rpe elevation", "neovascularization"]
elif label_name == "DME": keywords = ["intraretinal cysts", "thickening", "edema"]
elif label_name == "DRUSEN": keywords = ["rpe deposits", "drusen"]
else: keywords = ["normal foveal contour", "intact rpe"]
# Realistic mock bounding boxes for demonstration where pathology usually exists
box = [[0, 0], [0, 0]]
if label_name != "NORMAL":
box = [[80, 80], [150, 150]] # Center of the macula where fluid usually is
ground_truth[filename] = {
"label": label_name,
"box": box,
"keywords": keywords
}
counts[label_name] += 1
print(f"Downloaded {filename}...")
# Break if we have exactly 10 of each
if all(c >= target_per_class for c in counts.values()):
break
# Save the strictly formatted JSON
with open("ground_truth.json", "w") as f:
json.dump(ground_truth, f, indent=4)
print(f"\nSuccessfully downloaded {sum(counts.values())} real images!")
print("Auto-generated the new ground_truth.json answer key.")
print("You can now safely delete the 3 old 'sample_X.jpg' black images.")
if __name__ == "__main__":
fetch_real_oct_images()
|