-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathgenerate_synthetic_spheres.py
59 lines (43 loc) · 1.91 KB
/
generate_synthetic_spheres.py
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
import numpy as np
import cv2
import os
import pandas as pd
from numpy.random import default_rng
import argparse
if not os.path.exists('images'):
os.makedirs('images')
rng = default_rng()
parser = argparse.ArgumentParser()
parser.add_argument('--num_images', type=int, default=100, help='Number of images')
parser.add_argument('--image_size', type=int, default=96, help='Size of the images')
parser.add_argument('--radius_low', type=int, default=10, help='Minimum radius of the spheres')
parser.add_argument('--radius_high', type=int, default=20, help='Maximum radius of the spheres')
args = parser.parse_args()
# Define variables
num_images = args.num_images
image_size = args.image_size
radius_low = args.radius_low
radius_high = args.radius_high
center_loc = image_size // 2
center_scale = image_size // 4
labels = []
for i in range(num_images):
img = np.zeros((image_size, image_size), np.uint8)
# Adjust sphere center generation to use a Gaussian distribution centered around the image origin
u = int(rng.normal(loc=center_loc, scale=center_scale))
v = int(rng.normal(loc=center_loc, scale=center_scale))
# Clip the values to avoid the sphere going out of the image frame
u = np.clip(u, radius_high, image_size - radius_high)
v = np.clip(v, radius_high, image_size - radius_high)
# Make sure the radius is within the limit so the whole sphere is within the image frame
r = rng.integers(radius_low, min(u, v, image_size-u, image_size-v))
cv2.circle(img, (u, v), r, (255), -1)
noise = rng.normal(0.0, 0.8, img.shape).astype(np.uint8)
img = cv2.add(img, noise)
# Apply Gaussian blur for smoothness
img = cv2.GaussianBlur(img, (5, 5), 50)
img_path = f'images/img_{i}.png'
cv2.imwrite(img_path, img)
labels.append([img_path, u, v, r])
labels_df = pd.DataFrame(labels, columns=['image_path', 'center_u', 'center_v', 'radius'])
labels_df.to_csv('labels.csv', index=False)