476 lines
19 KiB
Python
476 lines
19 KiB
Python
import os
|
|
import random
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFilter
|
|
|
|
import json
|
|
|
|
# Define expanded classes (28 classes)
|
|
DEFAULT_CLASSES = [
|
|
"apfel_elstar",
|
|
"apfel_gala",
|
|
"apfel_granny_smith",
|
|
"banane_chiquita",
|
|
"banane_bio",
|
|
"birne",
|
|
"orange",
|
|
"zitrone",
|
|
"limette",
|
|
"erdbeere",
|
|
"blaubeere",
|
|
"weintraube_hell",
|
|
"weintraube_dunkel",
|
|
"pfirsich",
|
|
"tomate",
|
|
"gurke",
|
|
"kartoffel",
|
|
"karotte",
|
|
"zwiebel_gelb",
|
|
"zwiebel_rot",
|
|
"knoblauch",
|
|
"brokkoli",
|
|
"paprika_rot",
|
|
"paprika_gelb",
|
|
"paprika_gruen",
|
|
"champignon",
|
|
"zucchini",
|
|
"avocado"
|
|
]
|
|
|
|
def get_classes_from_db():
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
articles_path = os.path.join(base_dir, "articles.json")
|
|
if os.path.exists(articles_path):
|
|
try:
|
|
with open(articles_path, "r") as f:
|
|
articles = json.load(f)
|
|
keys = list(articles.keys())
|
|
if keys:
|
|
return keys
|
|
except Exception:
|
|
pass
|
|
return DEFAULT_CLASSES.copy()
|
|
|
|
CLASSES = get_classes_from_db()
|
|
|
|
def reload_classes():
|
|
new_classes = get_classes_from_db()
|
|
CLASSES.clear()
|
|
CLASSES.extend(new_classes)
|
|
print(f"Reloaded CLASSES: {len(CLASSES)} classes in memory.")
|
|
|
|
def draw_apple_elstar(draw):
|
|
draw.ellipse([45, 50, 155, 160], fill=(210, 45, 45, 255))
|
|
draw.ellipse([80, 45, 120, 60], fill=(210, 45, 45, 255))
|
|
draw.ellipse([55, 70, 110, 140], fill=(230, 140, 40, 180))
|
|
draw.line([(100, 55), (105, 25)], fill=(101, 67, 33, 255), width=4)
|
|
draw.ellipse([93, 50, 107, 57], fill=(60, 20, 20, 255))
|
|
|
|
def draw_apple_gala(draw):
|
|
draw.ellipse([45, 50, 155, 160], fill=(240, 190, 60, 255))
|
|
draw.ellipse([80, 45, 120, 60], fill=(240, 190, 60, 255))
|
|
for offset in range(-35, 45, 15):
|
|
draw.arc([45 + offset, 50, 155 + offset, 160], start=30, end=150, fill=(210, 40, 40, 200), width=4)
|
|
draw.arc([45 - offset, 50, 155 - offset, 160], start=30, end=150, fill=(210, 40, 40, 200), width=4)
|
|
draw.line([(100, 55), (103, 25)], fill=(101, 67, 33, 255), width=4)
|
|
draw.ellipse([93, 50, 107, 57], fill=(70, 50, 20, 255))
|
|
|
|
def draw_apple_granny_smith(draw):
|
|
# Bright green apple
|
|
draw.ellipse([45, 50, 155, 160], fill=(120, 200, 40, 255))
|
|
draw.ellipse([80, 45, 120, 60], fill=(120, 200, 40, 255))
|
|
# Light yellow highlight
|
|
draw.ellipse([60, 70, 95, 110], fill=(160, 230, 80, 180))
|
|
# Stem
|
|
draw.line([(100, 55), (105, 25)], fill=(101, 67, 33, 255), width=4)
|
|
draw.ellipse([93, 50, 107, 57], fill=(40, 70, 15, 255))
|
|
|
|
def draw_banana_chiquita(draw):
|
|
banana_pts = [
|
|
(40, 130), (50, 100), (70, 75), (100, 60), (130, 60),
|
|
(160, 75), (170, 95), (160, 95), (135, 80), (110, 80),
|
|
(85, 90), (65, 110), (50, 140)
|
|
]
|
|
draw.polygon(banana_pts, fill=(245, 215, 50, 255))
|
|
draw.polygon([(40, 130), (50, 140), (45, 145), (35, 135)], fill=(75, 100, 30, 255))
|
|
draw.polygon([(160, 75), (170, 95), (175, 90), (165, 70)], fill=(90, 75, 30, 255))
|
|
draw.ellipse([100, 65, 112, 75], fill=(10, 50, 150, 255))
|
|
draw.ellipse([103, 68, 109, 72], fill=(245, 215, 50, 255))
|
|
|
|
def draw_banana_bio(draw):
|
|
banana_pts = [
|
|
(45, 125), (55, 100), (75, 80), (100, 70), (125, 70),
|
|
(150, 82), (160, 98), (152, 98), (130, 88), (110, 85),
|
|
(88, 93), (70, 110), (53, 133)
|
|
]
|
|
draw.polygon(banana_pts, fill=(225, 220, 60, 255))
|
|
draw.polygon([(45, 125), (53, 133), (48, 138), (40, 130)], fill=(50, 90, 20, 255))
|
|
draw.polygon([(150, 82), (160, 98), (163, 93), (153, 77)], fill=(70, 90, 20, 255))
|
|
spots = [(70, 105), (85, 92), (100, 83), (115, 83), (130, 86)]
|
|
for sx, sy in spots:
|
|
draw.ellipse([sx-2, sy-2, sx+2, sy+2], fill=(80, 50, 20, 200))
|
|
draw.polygon([(95, 70), (105, 70), (102, 87), (92, 86)], fill=(120, 140, 90, 255))
|
|
|
|
def draw_pear(draw):
|
|
# Birne: pear shape (narrow top, wide bottom)
|
|
draw.ellipse([65, 50, 135, 110], fill=(180, 190, 50, 255)) # top
|
|
draw.ellipse([50, 90, 150, 170], fill=(180, 190, 50, 255)) # bottom
|
|
draw.ellipse([60, 75, 140, 145], fill=(180, 190, 50, 255)) # mid fill
|
|
# Orange-red blush
|
|
draw.ellipse([95, 100, 140, 150], fill=(210, 110, 30, 120))
|
|
# Stem
|
|
draw.line([(100, 55), (105, 25)], fill=(101, 67, 33, 255), width=3)
|
|
|
|
def draw_orange(draw):
|
|
draw.ellipse([45, 45, 155, 155], fill=(245, 120, 20, 255))
|
|
# Peel texture (little dots)
|
|
for _ in range(15):
|
|
tx = random.randint(60, 140)
|
|
ty = random.randint(60, 140)
|
|
draw.ellipse([tx, ty, tx+2, ty+2], fill=(215, 95, 5, 200))
|
|
|
|
def draw_lemon(draw):
|
|
# Zitrone: yellow oval with pointed ends
|
|
draw.ellipse([50, 65, 150, 135], fill=(245, 225, 40, 255))
|
|
draw.polygon([(40, 100), (52, 90), (52, 110)], fill=(245, 225, 40, 255)) # left tip
|
|
draw.polygon([(160, 100), (148, 90), (148, 110)], fill=(245, 225, 40, 255)) # right tip
|
|
|
|
def draw_lime(draw):
|
|
# Limette: green oval with pointed ends
|
|
draw.ellipse([55, 70, 145, 130], fill=(60, 160, 30, 255))
|
|
draw.polygon([(47, 100), (57, 92), (57, 108)], fill=(60, 160, 30, 255))
|
|
draw.polygon([(153, 100), (143, 92), (143, 108)], fill=(60, 160, 30, 255))
|
|
|
|
def draw_strawberry(draw):
|
|
# Erdbeere: heart shape
|
|
draw.polygon([(100, 165), (55, 80), (100, 60), (145, 80)], fill=(225, 30, 50, 255))
|
|
draw.ellipse([55, 70, 105, 100], fill=(225, 30, 50, 255))
|
|
draw.ellipse([95, 70, 145, 100], fill=(225, 30, 50, 255))
|
|
# Seeds (yellow dots)
|
|
for _ in range(20):
|
|
sx = random.randint(65, 135)
|
|
sy = random.randint(75, 140)
|
|
draw.ellipse([sx, sy, sx+2, sy+2], fill=(240, 220, 100, 255))
|
|
# Leafy top
|
|
draw.polygon([(100, 65), (85, 45), (95, 60), (100, 40), (105, 60), (115, 45)], fill=(40, 160, 50, 255))
|
|
|
|
def draw_blueberry(draw):
|
|
# Blaubeere: small blue circle
|
|
draw.ellipse([70, 70, 130, 130], fill=(45, 80, 160, 255))
|
|
# Crown (dark crown tip)
|
|
draw.polygon([(90, 70), (100, 62), (110, 70), (105, 75), (95, 75)], fill=(20, 45, 95, 255))
|
|
|
|
def draw_grape_green(draw):
|
|
# Weintraube hell: cluster of small green circles
|
|
coords = [
|
|
(100, 60), (85, 75), (115, 75), (70, 95), (100, 95), (130, 95),
|
|
(85, 115), (115, 115), (100, 135), (100, 155)
|
|
]
|
|
# Stem
|
|
draw.line([(100, 60), (100, 35)], fill=(110, 140, 50, 255), width=3)
|
|
for cx, cy in coords:
|
|
draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=(160, 210, 80, 255))
|
|
draw.ellipse([cx-16, cy-16, cx+16, cy+16], outline=(130, 180, 60, 255), width=1)
|
|
|
|
def draw_grape_dark(draw):
|
|
# Weintraube dunkel: cluster of small dark purple circles
|
|
coords = [
|
|
(100, 60), (85, 75), (115, 75), (70, 95), (100, 95), (130, 95),
|
|
(85, 115), (115, 115), (100, 135), (100, 155)
|
|
]
|
|
draw.line([(100, 60), (100, 35)], fill=(90, 70, 110, 255), width=3)
|
|
for cx, cy in coords:
|
|
draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=(65, 30, 105, 255))
|
|
draw.ellipse([cx-16, cy-16, cx+16, cy+16], outline=(45, 15, 80, 255), width=1)
|
|
|
|
def draw_peach(draw):
|
|
# Pfirsich: fuzzy orange-pink circle
|
|
draw.ellipse([45, 50, 155, 160], fill=(245, 130, 80, 255))
|
|
draw.ellipse([47, 50, 120, 160], fill=(245, 80, 110, 255)) # pink cheek
|
|
# Indentation line
|
|
draw.arc([40, 50, 150, 160], start=270, end=90, fill=(185, 45, 45, 255), width=2)
|
|
# stem
|
|
draw.line([(100, 53), (102, 30)], fill=(90, 70, 40, 255), width=3)
|
|
|
|
def draw_tomato(draw):
|
|
draw.ellipse([50, 55, 150, 155], fill=(225, 30, 30, 255))
|
|
draw.ellipse([65, 70, 85, 85], fill=(255, 255, 255, 180))
|
|
cx, cy = 100, 55
|
|
stem_pts = [
|
|
(cx, cy), (cx-15, cy-10), (cx-5, cy-2),
|
|
(cx, cy-20), (cx+5, cy-2), (cx+15, cy-10),
|
|
(cx+8, cy+5), (cx, cy+8), (cx-8, cy+5)
|
|
]
|
|
draw.polygon(stem_pts, fill=(34, 139, 34, 255))
|
|
draw.line([(cx, cy), (cx-5, cy-25)], fill=(34, 139, 34, 255), width=4)
|
|
|
|
def draw_cucumber(draw):
|
|
draw.rounded_rectangle([35, 80, 165, 120], radius=20, fill=(20, 100, 40, 255))
|
|
for x in range(50, 150, 15):
|
|
draw.line([(x, 85), (x+5, 115)], fill=(30, 130, 55, 255), width=2)
|
|
draw.ellipse([32, 95, 38, 105], fill=(139, 115, 85, 255))
|
|
draw.ellipse([162, 95, 168, 105], fill=(100, 80, 40, 255))
|
|
|
|
def draw_potato(draw):
|
|
# Kartoffel: irregular brown-yellow oval
|
|
draw.rounded_rectangle([45, 65, 155, 135], radius=35, fill=(190, 160, 95, 255))
|
|
# Add potato eyes (dark brown dots)
|
|
eyes = [(65, 80), (135, 90), (115, 120), (80, 125), (100, 85)]
|
|
for ex, ey in eyes:
|
|
draw.ellipse([ex-2, ey-1, ex+2, ey+1], fill=(110, 85, 45, 255))
|
|
draw.arc([ex-3, ey-3, ex+3, ey+3], start=0, end=180, fill=(110, 85, 45, 255), width=1)
|
|
|
|
def draw_carrot(draw):
|
|
# Karotte: tapered orange shape
|
|
draw.polygon([(40, 85), (160, 100), (40, 115)], fill=(245, 110, 20, 255))
|
|
draw.ellipse([35, 85, 45, 115], fill=(245, 110, 20, 255))
|
|
# Green leaves on one end (left)
|
|
draw.polygon([(38, 100), (15, 80), (30, 95), (10, 100), (30, 105), (15, 120)], fill=(45, 155, 55, 255))
|
|
# Texture lines
|
|
for x in range(60, 140, 20):
|
|
draw.line([(x, 92), (x, 108)], fill=(215, 85, 10, 255), width=2)
|
|
|
|
def draw_onion_yellow(draw):
|
|
# Zwiebel gelb: golden bulb shape
|
|
draw.ellipse([50, 60, 150, 150], fill=(195, 145, 75, 255))
|
|
draw.polygon([(100, 40), (80, 65), (120, 65)], fill=(195, 145, 75, 255)) # top tip
|
|
# Root hairs at bottom
|
|
draw.line([(90, 150), (85, 162)], fill=(230, 210, 170, 255), width=2)
|
|
draw.line([(100, 150), (100, 165)], fill=(230, 210, 170, 255), width=2)
|
|
draw.line([(110, 150), (115, 162)], fill=(230, 210, 170, 255), width=2)
|
|
|
|
def draw_onion_red(draw):
|
|
# Zwiebel rot: red-purple bulb shape
|
|
draw.ellipse([50, 60, 150, 150], fill=(140, 35, 95, 255))
|
|
draw.polygon([(100, 40), (80, 65), (120, 65)], fill=(140, 35, 95, 255))
|
|
# Light stripes
|
|
draw.arc([55, 60, 145, 150], start=180, end=360, fill=(185, 75, 135, 255), width=2)
|
|
# Root hairs
|
|
draw.line([(95, 150), (95, 162)], fill=(210, 190, 180, 255), width=2)
|
|
draw.line([(105, 150), (105, 162)], fill=(210, 190, 180, 255), width=2)
|
|
|
|
def draw_garlic(draw):
|
|
# Knoblauch: off-white bulb with segment lines
|
|
draw.ellipse([50, 60, 150, 150], fill=(235, 230, 220, 255))
|
|
draw.polygon([(100, 42), (85, 65), (115, 65)], fill=(235, 230, 220, 255))
|
|
# Segment lines
|
|
draw.arc([65, 60, 135, 150], start=270, end=90, fill=(195, 190, 180, 255), width=2)
|
|
draw.arc([80, 60, 120, 150], start=270, end=90, fill=(195, 190, 180, 255), width=2)
|
|
# Roots
|
|
draw.line([(100, 150), (100, 160)], fill=(160, 150, 130, 255), width=2)
|
|
|
|
def draw_broccoli(draw):
|
|
# Brokkoli: green tree-like structure
|
|
# Stalk
|
|
draw.rounded_rectangle([80, 110, 120, 170], radius=8, fill=(110, 175, 75, 255))
|
|
# Florets (overlapping circles)
|
|
draw.ellipse([60, 60, 110, 110], fill=(30, 110, 50, 255))
|
|
draw.ellipse([90, 50, 150, 100], fill=(30, 110, 50, 255))
|
|
draw.ellipse([75, 80, 135, 130], fill=(30, 110, 50, 255))
|
|
# Add tiny dots/texture
|
|
for _ in range(30):
|
|
tx = random.randint(65, 135)
|
|
ty = random.randint(55, 120)
|
|
draw.ellipse([tx, ty, tx+2, ty+2], fill=(60, 150, 80, 255))
|
|
|
|
def draw_pepper_red(draw):
|
|
# Paprika rot: blocky rounded shape
|
|
draw.rounded_rectangle([50, 60, 150, 150], radius=25, fill=(210, 30, 30, 255))
|
|
# Lobes indentations
|
|
draw.arc([45, 60, 155, 150], start=230, end=310, fill=(160, 15, 15, 255), width=3)
|
|
# Green stem
|
|
draw.line([(100, 60), (100, 35)], fill=(45, 130, 45, 255), width=5)
|
|
|
|
def draw_pepper_yellow(draw):
|
|
# Paprika gelb: blocky rounded shape
|
|
draw.rounded_rectangle([50, 60, 150, 150], radius=25, fill=(240, 190, 20, 255))
|
|
draw.arc([45, 60, 155, 150], start=230, end=310, fill=(195, 145, 10, 255), width=3)
|
|
draw.line([(100, 60), (100, 35)], fill=(45, 130, 45, 255), width=5)
|
|
|
|
def draw_pepper_green(draw):
|
|
# Paprika grün: blocky rounded shape
|
|
draw.rounded_rectangle([50, 60, 150, 150], radius=25, fill=(35, 115, 45, 255))
|
|
draw.arc([45, 60, 155, 150], start=230, end=310, fill=(20, 85, 25, 255), width=3)
|
|
draw.line([(100, 60), (100, 35)], fill=(25, 80, 25, 255), width=5)
|
|
|
|
def draw_mushroom(draw):
|
|
# Champignon: white-brown stalk and cap
|
|
# Stalk
|
|
draw.rounded_rectangle([85, 100, 115, 160], radius=6, fill=(225, 215, 205, 255))
|
|
# Cap (umbrella)
|
|
draw.ellipse([45, 60, 155, 120], fill=(210, 200, 190, 255))
|
|
draw.rectangle([45, 90, 155, 120], fill=(210, 200, 190, 255)) # bottom flat
|
|
# Gills (brown under cap)
|
|
draw.line([(45, 120), (155, 120)], fill=(120, 100, 90, 255), width=3)
|
|
|
|
def draw_zucchini(draw):
|
|
# Zucchini: long cylinder, thicker than cucumber, dark green
|
|
draw.rounded_rectangle([30, 75, 170, 125], radius=22, fill=(15, 75, 35, 255))
|
|
# Speckled light green stripes
|
|
for x in range(45, 155, 20):
|
|
draw.line([(x, 80), (x+8, 120)], fill=(35, 115, 60, 180), width=3)
|
|
# Thick stem
|
|
draw.ellipse([167, 95, 173, 105], fill=(80, 95, 50, 255))
|
|
|
|
def draw_avocado(draw):
|
|
# Halved avocado: pear outline, yellow-green center, brown seed
|
|
# Outer skin
|
|
draw.ellipse([60, 50, 140, 120], fill=(15, 45, 20, 255))
|
|
draw.ellipse([45, 85, 155, 170], fill=(15, 45, 20, 255))
|
|
# Light green inner meat
|
|
draw.ellipse([63, 53, 137, 117], fill=(185, 215, 110, 255))
|
|
draw.ellipse([49, 89, 151, 166], fill=(185, 215, 110, 255))
|
|
# Yellow center
|
|
draw.ellipse([58, 98, 142, 158], fill=(215, 235, 135, 255))
|
|
# Brown pit (seed) in center
|
|
draw.ellipse([78, 105, 122, 149], fill=(110, 65, 30, 255))
|
|
draw.ellipse([82, 109, 100, 125], fill=(255, 255, 255, 80)) # Pit reflection
|
|
|
|
def create_object_image(class_idx, angle, scale):
|
|
obj_img = Image.new("RGBA", (200, 200), (0, 0, 0, 0))
|
|
draw = ImageDraw.Draw(obj_img)
|
|
|
|
# Map class index to drawing functions
|
|
draw_funcs = [
|
|
draw_apple_elstar, draw_apple_gala, draw_apple_granny_smith,
|
|
draw_banana_chiquita, draw_banana_bio, draw_pear,
|
|
draw_orange, draw_lemon, draw_lime,
|
|
draw_strawberry, draw_blueberry, draw_grape_green,
|
|
draw_grape_dark, draw_peach, draw_tomato,
|
|
draw_cucumber, draw_potato, draw_carrot,
|
|
draw_onion_yellow, draw_onion_red, draw_garlic,
|
|
draw_broccoli, draw_pepper_red, draw_pepper_yellow,
|
|
draw_pepper_green, draw_mushroom, draw_zucchini,
|
|
draw_avocado
|
|
]
|
|
|
|
if class_idx < len(draw_funcs):
|
|
draw_funcs[class_idx](draw)
|
|
|
|
rotated = obj_img.rotate(angle, resample=Image.Resampling.BILINEAR, expand=True)
|
|
|
|
if scale != 1.0:
|
|
new_w = int(rotated.width * scale)
|
|
new_h = int(rotated.height * scale)
|
|
rotated = rotated.resize((new_w, new_h), resample=Image.Resampling.BILINEAR)
|
|
|
|
bbox = rotated.getbbox()
|
|
if bbox is None:
|
|
return None, None
|
|
|
|
cropped = rotated.crop(bbox)
|
|
return cropped, bbox
|
|
|
|
def generate_background(width, height):
|
|
bg = Image.new("RGB", (width, height), (210, 215, 220))
|
|
draw = ImageDraw.Draw(bg)
|
|
grid_color = (185, 190, 195)
|
|
for x in range(0, width, 80):
|
|
draw.line([(x, 0), (x, height)], fill=grid_color, width=2)
|
|
for y in range(0, height, 80):
|
|
draw.line([(0, y), (width, y)], fill=grid_color, width=2)
|
|
draw.rectangle([10, 10, width-10, height-10], outline=(170, 175, 180), width=3)
|
|
|
|
arr = np.array(bg, dtype=np.float32)
|
|
x_grad = np.linspace(-15, 15, width)
|
|
y_grad = np.linspace(-15, 15, height)
|
|
xx, yy = np.meshgrid(x_grad, y_grad)
|
|
arr[:, :, 0] += xx + yy
|
|
arr[:, :, 1] += xx + yy
|
|
arr[:, :, 2] += xx + yy
|
|
noise = np.random.normal(0, 3, (height, width, 3))
|
|
arr += noise
|
|
arr = np.clip(arr, 0, 255).astype(np.uint8)
|
|
return Image.fromarray(arr)
|
|
|
|
def generate_sample(width=640, height=480):
|
|
bg = generate_background(width, height)
|
|
num_objects = random.randint(1, 3)
|
|
annotations = []
|
|
placed_boxes = []
|
|
|
|
for _ in range(num_objects):
|
|
class_idx = random.randint(0, len(CLASSES)-1)
|
|
angle = random.randint(0, 360)
|
|
scale = random.uniform(0.7, 1.2)
|
|
|
|
obj_img, _ = create_object_image(class_idx, angle, scale)
|
|
if obj_img is None:
|
|
continue
|
|
|
|
ow, oh = obj_img.size
|
|
placed = False
|
|
|
|
for attempt in range(25):
|
|
x = random.randint(30, width - ow - 30)
|
|
y = random.randint(30, height - oh - 30)
|
|
|
|
overlap = False
|
|
for px1, py1, px2, py2 in placed_boxes:
|
|
ix1 = max(x, px1)
|
|
iy1 = max(y, py1)
|
|
ix2 = min(x + ow, px2)
|
|
iy2 = min(y + oh, py2)
|
|
|
|
if ix2 > ix1 and iy2 > iy1:
|
|
inter_area = (ix2 - ix1) * (iy2 - iy1)
|
|
box_area = ow * oh
|
|
p_area = (px2 - px1) * (py2 - py1)
|
|
iou = inter_area / float(box_area + p_area - inter_area)
|
|
if iou > 0.12:
|
|
overlap = True
|
|
break
|
|
|
|
if not overlap:
|
|
bg.paste(obj_img, (x, y), obj_img)
|
|
placed_boxes.append((x, y, x + ow, y + oh))
|
|
x_center = (x + ow / 2.0) / width
|
|
y_center = (y + oh / 2.0) / height
|
|
w_norm = ow / float(width)
|
|
h_norm = oh / float(height)
|
|
annotations.append(f"{class_idx} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
|
|
placed = True
|
|
break
|
|
|
|
if not placed and len(placed_boxes) == 0:
|
|
x = random.randint(30, width - ow - 30)
|
|
y = random.randint(30, height - oh - 30)
|
|
bg.paste(obj_img, (x, y), obj_img)
|
|
placed_boxes.append((x, y, x + ow, y + oh))
|
|
x_center = (x + ow / 2.0) / width
|
|
y_center = (y + oh / 2.0) / height
|
|
w_norm = ow / float(width)
|
|
h_norm = oh / float(height)
|
|
annotations.append(f"{class_idx} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
|
|
|
|
bg = bg.filter(ImageFilter.GaussianBlur(0.5))
|
|
return bg, annotations
|
|
|
|
def build_dataset(base_dir="dataset", train_count=200, val_count=50):
|
|
os.makedirs(os.path.join(base_dir, "images", "train"), exist_ok=True)
|
|
os.makedirs(os.path.join(base_dir, "images", "val"), exist_ok=True)
|
|
os.makedirs(os.path.join(base_dir, "labels", "train"), exist_ok=True)
|
|
os.makedirs(os.path.join(base_dir, "labels", "val"), exist_ok=True)
|
|
|
|
with open(os.path.join(base_dir, "classes.txt"), "w") as f:
|
|
f.write("\n".join(CLASSES))
|
|
|
|
print(f"Generating {train_count} training samples for {len(CLASSES)} classes...")
|
|
for idx in range(train_count):
|
|
img, ann = generate_sample()
|
|
img.save(os.path.join(base_dir, "images", "train", f"train_{idx:04d}.jpg"), quality=90)
|
|
with open(os.path.join(base_dir, "labels", "train", f"train_{idx:04d}.txt"), "w") as f:
|
|
f.write("\n".join(ann))
|
|
|
|
print(f"Generating {val_count} validation samples...")
|
|
for idx in range(val_count):
|
|
img, ann = generate_sample()
|
|
img.save(os.path.join(base_dir, "images", "val", f"val_{idx:04d}.jpg"), quality=90)
|
|
with open(os.path.join(base_dir, "labels", "val", f"val_{idx:04d}.txt"), "w") as f:
|
|
f.write("\n".join(ann))
|
|
|
|
print("Dataset generation complete!")
|
|
|
|
if __name__ == "__main__":
|
|
build_dataset(train_count=200, val_count=50)
|