Spaces:
Sleeping
Sleeping
File size: 4,588 Bytes
559e904 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import gzip
import pickle
from pathlib import Path
import numpy as np
from scipy.interpolate import interp1d
from scipy.ndimage import binary_opening
from tqdm import tqdm
def raw_to_normalized_features(
points, num_points=256, jump_threshold=0.3, structure_size=3
):
# 1. Первичная нормализация и маска (Bounding Box)
x_raw, y_raw = points[:, 0], points[:, 1]
x_min, x_max = np.min(x_raw), np.max(x_raw)
y_min, y_max = np.min(y_raw), np.max(y_raw)
x_range = x_max - x_min if x_max != x_min else 1.0
y_range = y_max - y_min if y_max != y_min else 1.0
x_norm = (x_raw - x_min) / x_range
y_norm = np.full_like(y_raw, 0.5) if y_max == y_min else (y_raw - y_min) / y_range
dx = np.diff(x_norm)
gap_indices = np.where(dx > 3 * np.median(dx))[0] if len(dx) > 0 else []
valid_intervals = []
start_idx = 0
for gap_idx in gap_indices:
valid_intervals.append((x_norm[start_idx], x_norm[gap_idx]))
start_idx = gap_idx + 1
valid_intervals.append((x_norm[start_idx], x_norm[-1]))
x_fixed = np.linspace(0, 1, num_points)
sort_idx = np.argsort(x_norm)
y_fixed = np.interp(x_fixed, x_norm[sort_idx], y_norm[sort_idx])
mask = np.zeros_like(x_fixed, dtype=np.float32)
for start, end in valid_intervals:
mask[(x_fixed >= start) & (x_fixed <= end)] = 1.0
# 2. Удаление ложных асимптот по оси Y
valid_y_orig = y_fixed[mask == 1]
if len(valid_y_orig) < 2:
return None
y_range_mask = np.nanmax(valid_y_orig) - np.nanmin(valid_y_orig)
if y_range_mask > 1e-5:
y_diff = np.abs(np.diff(y_fixed))
asymptote_indices = np.where(
(y_diff > y_range_mask * jump_threshold)
& (mask[:-1] == 1)
& (mask[1:] == 1)
)[0]
for idx in asymptote_indices:
mask[idx] = 0
# 3. Удаление ошметков (< 3 пикселей)
structure = np.ones(structure_size)
cleaned_mask = binary_opening(mask, structure=structure).astype(int) # type: ignore
if np.sum(cleaned_mask) < 5:
return None
# 4. Обрезка BBox по X и финальная интерполяция
valid_indices = np.where(cleaned_mask == 1)[0]
i_min, i_max = valid_indices[0], valid_indices[-1]
if i_max == i_min:
return None
x_orig = np.linspace(0, 1, len(y_fixed))
x_min_bbox, x_max_bbox = x_orig[i_min], x_orig[i_max]
x_new = np.linspace(x_min_bbox, x_max_bbox, num_points)
y_interp_func = interp1d(
x_orig, y_fixed, kind="linear", bounds_error=False, fill_value=np.nan
)
y_resampled = y_interp_func(x_new)
mask_interp_func = interp1d(
x_orig, cleaned_mask, kind="nearest", bounds_error=False, fill_value=0
)
mask_resampled = mask_interp_func(x_new).astype(int)
# 5. Финальный Min-Max Scaling по Y
valid_y_resampled = y_resampled[mask_resampled == 1]
if len(valid_y_resampled) == 0 or np.isnan(valid_y_resampled).all():
return None
y_min_val, y_max_val = np.nanmin(valid_y_resampled), np.nanmax(valid_y_resampled)
diff = y_max_val - y_min_val
y_scaled = (
np.full_like(y_resampled, 0.5)
if diff < 1e-6
else (y_resampled - y_min_val) / diff
)
y_scaled[mask_resampled == 0] = 0.0
return np.vstack((y_scaled.astype(np.float32), mask_resampled.astype(np.float32)))
def run_normalization(input_dir: Path, output_dir: Path):
output_dir.mkdir(parents=True, exist_ok=True)
input_files = sorted(input_dir.glob("chunk_*.pkl.gz"))
total_processed = 0
for filepath in tqdm(input_files, desc="Нормализация и чистка графиков"):
with gzip.open(filepath, "rb") as f:
data = pickle.load(f)
new_data = []
for item in data:
features = raw_to_normalized_features(item["points"])
if features is not None:
new_data.append(
{
"expr_str": item["expr_str"],
"expr_instantiated_str": item["expr_instantiated_str"],
"tokens": item["tokens"],
"features": features,
}
)
if new_data:
with gzip.open(output_dir / filepath.name, "wb") as f_out:
pickle.dump(new_data, f_out)
total_processed += len(new_data)
print(f"Идеальных графиков сохранено: {total_processed}")
|