251 lines
8.6 KiB
Python
251 lines
8.6 KiB
Python
|
import os
|
|||
|
|
import re
|
|||
|
|
import random
|
|||
|
|
from collections import defaultdict
|
|||
|
|
from typing import Dict, List, Optional, Set, Any
|
|||
|
|
|
|||
|
|
from nets import nn
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _image_to_label_path(img_path: str) -> str:
|
|||
|
|
"""
|
|||
|
|
Convert an image path like ".../images/train2017/xxx.jpg"
|
|||
|
|
to the corresponding label path ".../labels/train2017/xxx.txt".
|
|||
|
|
Works for POSIX/Windows separators.
|
|||
|
|
"""
|
|||
|
|
# swap "/images/" (or "\images\") to "/labels/"
|
|||
|
|
label_path = re.sub(r"([/\\])images([/\\])", r"\1labels\2", img_path)
|
|||
|
|
# swap extension to .txt
|
|||
|
|
root, _ = os.path.splitext(label_path)
|
|||
|
|
return root + ".txt"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _parse_yolo_label_file(label_path: str) -> Set[int]:
|
|||
|
|
"""
|
|||
|
|
Return a set of class_ids found in a YOLO .txt label file.
|
|||
|
|
Empty file -> empty set. Missing file -> empty set.
|
|||
|
|
Robust to blank lines / trailing spaces.
|
|||
|
|
"""
|
|||
|
|
class_ids: Set[int] = set()
|
|||
|
|
if not os.path.exists(label_path):
|
|||
|
|
return class_ids
|
|||
|
|
try:
|
|||
|
|
with open(label_path, "r", encoding="utf-8") as f:
|
|||
|
|
for line in f:
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
# YOLO format: cls cx cy w h
|
|||
|
|
parts = line.split()
|
|||
|
|
if not parts:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
cls = int(parts[0])
|
|||
|
|
except ValueError:
|
|||
|
|
# handle weird case like '23.0'
|
|||
|
|
try:
|
|||
|
|
cls = int(float(parts[0]))
|
|||
|
|
except ValueError:
|
|||
|
|
# skip malformed line
|
|||
|
|
continue
|
|||
|
|
class_ids.add(cls)
|
|||
|
|
except Exception:
|
|||
|
|
# If the file can't be read for some reason, treat as no labels
|
|||
|
|
return set()
|
|||
|
|
return class_ids
|
|||
|
|
|
|||
|
|
|
|||
|
|
def divide_trainset(
|
|||
|
|
trainset_path: str,
|
|||
|
|
num_local_class: int,
|
|||
|
|
num_client: int,
|
|||
|
|
min_data: int,
|
|||
|
|
max_data: int,
|
|||
|
|
mode: str = "overlap", # "overlap" or "disjoint"
|
|||
|
|
seed: Optional[int] = None,
|
|||
|
|
) -> Dict[str, Any]:
|
|||
|
|
"""
|
|||
|
|
Build a federated split from a YOLO dataset list file.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
trainset_path: path to a .txt file containing one image path per line
|
|||
|
|
e.g. /COCO/images/train2017/1111.jpg
|
|||
|
|
num_local_class: how many distinct classes to sample for each client
|
|||
|
|
num_client: number of clients
|
|||
|
|
min_data: minimum number of images per client
|
|||
|
|
max_data: maximum number of images per client
|
|||
|
|
mode: "overlap" -> images may be shared across clients
|
|||
|
|
"disjoint" -> each image is used by at most one client
|
|||
|
|
seed: optional random seed for reproducibility
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
trainset_divided = {
|
|||
|
|
"users": ["c_00001", ...],
|
|||
|
|
"user_data": {
|
|||
|
|
"c_00001": {"filename": [img_path, ...]},
|
|||
|
|
...
|
|||
|
|
},
|
|||
|
|
"num_samples": [len(list_for_user1), len(list_for_user2), ...]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
Example:
|
|||
|
|
dataset = divide_trainset(
|
|||
|
|
trainset_path="/COCO/train2017.txt",
|
|||
|
|
num_local_class=3,
|
|||
|
|
num_client=5,
|
|||
|
|
min_data=10,
|
|||
|
|
max_data=20,
|
|||
|
|
mode="disjoint", # or "overlap"
|
|||
|
|
seed=42
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print(dataset["users"]) # ['c_00001', ..., 'c_00005']
|
|||
|
|
print(dataset["num_samples"]) # e.g. [10, 12, 18, 9, 15]
|
|||
|
|
print(dataset["user_data"]["c_00001"]["filename"][:3])
|
|||
|
|
"""
|
|||
|
|
if seed is not None:
|
|||
|
|
random.seed(seed)
|
|||
|
|
|
|||
|
|
# ---- Basic validations (defensive programming) ----
|
|||
|
|
if num_client <= 0:
|
|||
|
|
raise ValueError("num_client must be > 0")
|
|||
|
|
if num_local_class <= 0:
|
|||
|
|
raise ValueError("num_local_class must be > 0")
|
|||
|
|
if min_data < 0 or max_data < 0:
|
|||
|
|
raise ValueError("min_data/max_data must be >= 0")
|
|||
|
|
if max_data < min_data:
|
|||
|
|
raise ValueError("max_data must be >= min_data")
|
|||
|
|
if mode not in {"overlap", "disjoint"}:
|
|||
|
|
raise ValueError('mode must be "overlap" or "disjoint"')
|
|||
|
|
|
|||
|
|
# ---- 1) Read image list ----
|
|||
|
|
with open(trainset_path, "r", encoding="utf-8") as f:
|
|||
|
|
all_images_raw = [ln.strip() for ln in f if ln.strip()]
|
|||
|
|
|
|||
|
|
# Normalize and deduplicate image paths (safe)
|
|||
|
|
all_images: List[str] = []
|
|||
|
|
seen = set()
|
|||
|
|
for p in all_images_raw:
|
|||
|
|
# keep exact string (don’t join with cwd), just normalize slashes
|
|||
|
|
norm = os.path.normpath(p)
|
|||
|
|
if norm not in seen:
|
|||
|
|
seen.add(norm)
|
|||
|
|
all_images.append(norm)
|
|||
|
|
|
|||
|
|
# ---- 2) Build mappings from labels ----
|
|||
|
|
class_to_images: Dict[int, Set[str]] = defaultdict(set)
|
|||
|
|
image_to_classes: Dict[str, Set[int]] = {}
|
|||
|
|
|
|||
|
|
missing_label_files = 0
|
|||
|
|
empty_label_files = 0
|
|||
|
|
parsed_images = 0
|
|||
|
|
|
|||
|
|
for img in all_images:
|
|||
|
|
lbl = _image_to_label_path(img)
|
|||
|
|
if not os.path.exists(lbl):
|
|||
|
|
# Missing labels: skip image (no class info)
|
|||
|
|
missing_label_files += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
classes = _parse_yolo_label_file(lbl)
|
|||
|
|
if not classes:
|
|||
|
|
# No objects in this image -> skip (no class bucket)
|
|||
|
|
empty_label_files += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
image_to_classes[img] = classes
|
|||
|
|
for c in classes:
|
|||
|
|
class_to_images[c].add(img)
|
|||
|
|
parsed_images += 1
|
|||
|
|
|
|||
|
|
if not class_to_images:
|
|||
|
|
# No usable images found
|
|||
|
|
return {
|
|||
|
|
"users": [f"c_{i + 1:05d}" for i in range(num_client)],
|
|||
|
|
"user_data": {f"c_{i + 1:05d}": {"filename": []} for i in range(num_client)},
|
|||
|
|
"num_samples": [0 for _ in range(num_client)],
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
all_classes: List[int] = sorted(class_to_images.keys())
|
|||
|
|
# Available pool for disjoint mode (only images with labels)
|
|||
|
|
available_images: Set[str] = set(image_to_classes.keys())
|
|||
|
|
|
|||
|
|
# ---- 3) Allocate to clients ----
|
|||
|
|
result = {"users": [], "user_data": {}, "num_samples": []}
|
|||
|
|
|
|||
|
|
for cid in range(num_client):
|
|||
|
|
user_id = f"c_{cid + 1:05d}"
|
|||
|
|
result["users"].append(user_id)
|
|||
|
|
|
|||
|
|
# Pick the classes for this client (sample without replacement from global class set)
|
|||
|
|
k = min(num_local_class, len(all_classes))
|
|||
|
|
chosen_classes = random.sample(all_classes, k) if k > 0 else []
|
|||
|
|
|
|||
|
|
# Decide how many samples for this client
|
|||
|
|
need = min_data if min_data == max_data else random.randint(min_data, max_data)
|
|||
|
|
|
|||
|
|
# Build the candidate pool for this client
|
|||
|
|
if mode == "overlap":
|
|||
|
|
pool_set: Set[str] = set()
|
|||
|
|
for c in chosen_classes:
|
|||
|
|
pool_set.update(class_to_images[c])
|
|||
|
|
else: # "disjoint": restrict to currently available images
|
|||
|
|
pool_set = set()
|
|||
|
|
for c in chosen_classes:
|
|||
|
|
# intersect with available images
|
|||
|
|
pool_set.update(class_to_images[c] & available_images)
|
|||
|
|
|
|||
|
|
# Deduplicate and sample
|
|||
|
|
pool_list = list(pool_set)
|
|||
|
|
if len(pool_list) <= need:
|
|||
|
|
chosen_imgs = pool_list[:] # take all (can be fewer than need)
|
|||
|
|
else:
|
|||
|
|
chosen_imgs = random.sample(pool_list, need)
|
|||
|
|
|
|||
|
|
# Record for the user
|
|||
|
|
result["user_data"][user_id] = {"filename": chosen_imgs}
|
|||
|
|
result["num_samples"].append(len(chosen_imgs))
|
|||
|
|
|
|||
|
|
# If disjoint, remove selected images from availability everywhere
|
|||
|
|
if mode == "disjoint" and chosen_imgs:
|
|||
|
|
for img in chosen_imgs:
|
|||
|
|
if img in available_images:
|
|||
|
|
available_images.remove(img)
|
|||
|
|
# remove from every class bucket this image belongs to
|
|||
|
|
for c in image_to_classes.get(img, []):
|
|||
|
|
if img in class_to_images[c]:
|
|||
|
|
class_to_images[c].remove(img)
|
|||
|
|
# Optional: prune empty classes from all_classes to speed up later loops
|
|||
|
|
# (keep list stable; just skip empties naturally)
|
|||
|
|
|
|||
|
|
# (Optional) You can print some quick diagnostics if helpful:
|
|||
|
|
# print(f"[INFO] Parsed images with labels: {parsed_images}")
|
|||
|
|
# print(f"[INFO] Missing label files: {missing_label_files}")
|
|||
|
|
# print(f"[INFO] Empty label files: {empty_label_files}")
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def init_model(model_name, num_classes):
|
|||
|
|
"""
|
|||
|
|
Initialize the model for a specific learning task
|
|||
|
|
Args:
|
|||
|
|
:param model_name: Name of the model
|
|||
|
|
:param num_classes: Number of classes
|
|||
|
|
"""
|
|||
|
|
model = None
|
|||
|
|
if model_name == "yolo_v11_n":
|
|||
|
|
model = nn.yolo_v11_n(num_classes=num_classes)
|
|||
|
|
elif model_name == "yolo_v11_s":
|
|||
|
|
model = nn.yolo_v11_s(num_classes=num_classes)
|
|||
|
|
elif model_name == "yolo_v11_m":
|
|||
|
|
model = nn.yolo_v11_m(num_classes=num_classes)
|
|||
|
|
elif model_name == "yolo_v11_l":
|
|||
|
|
model = nn.yolo_v11_l(num_classes=num_classes)
|
|||
|
|
elif model_name == "yolo_v11_x":
|
|||
|
|
model = nn.yolo_v11_x(num_classes=num_classes)
|
|||
|
|
else:
|
|||
|
|
raise ValueError("Model {} is not supported.".format(model_name))
|
|||
|
|
|
|||
|
|
return model
|