496 lines
21 KiB
Python
496 lines
21 KiB
Python
"""
|
||
=============================================================================
|
||
POLE ENGINEERING PROJECT AUTOMATION SCRIPT
|
||
Automatisation des projets d'ingénierie de poteaux utilitaires
|
||
=============================================================================
|
||
Workflow:
|
||
1. Create project folder structure
|
||
2. Read source data from "départ" folder
|
||
3. Rename & move folders to RSA-UDS with business logic (sd/d)
|
||
4. Identify pole/anchor images via Gemini AI
|
||
5. Convert images & Excel files to PDF
|
||
6. Move all PDFs to compilation folder
|
||
=============================================================================
|
||
REQUIREMENTS:
|
||
pip install pillow openpyxl reportlab google-generativeai
|
||
=============================================================================
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import sys
|
||
import shutil
|
||
import base64
|
||
import argparse
|
||
import logging
|
||
from pathlib import Path
|
||
|
||
# ── Third-party ──────────────────────────────────────────────────────────────
|
||
try:
|
||
from PIL import Image
|
||
except ImportError:
|
||
print("❌ Pillow missing. Run: pip install pillow")
|
||
sys.exit(1)
|
||
|
||
try:
|
||
import openpyxl
|
||
from openpyxl import load_workbook
|
||
except ImportError:
|
||
print("❌ openpyxl missing. Run: pip install openpyxl")
|
||
sys.exit(1)
|
||
|
||
try:
|
||
from reportlab.lib.pagesizes import letter, landscape
|
||
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer
|
||
from reportlab.lib.styles import getSampleStyleSheet
|
||
from reportlab.lib import colors
|
||
from reportlab.pdfgen import canvas
|
||
except ImportError:
|
||
print("❌ reportlab missing. Run: pip install reportlab")
|
||
sys.exit(1)
|
||
|
||
try:
|
||
import google.generativeai as genai
|
||
GEMINI_AVAILABLE = True
|
||
except ImportError:
|
||
GEMINI_AVAILABLE = False
|
||
print("⚠️ google-generativeai not installed. Gemini AI disabled.")
|
||
print(" Run: pip install google-generativeai")
|
||
|
||
# ── Logging ───────────────────────────────────────────────────────────────────
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s %(levelname)-8s %(message)s",
|
||
datefmt="%H:%M:%S",
|
||
)
|
||
log = logging.getLogger(__name__)
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# CONFIGURATION – edit these paths to match your environment
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
DEFAULT_BASE_DIR = Path.home() / "PycharmProjects" / "structure" / "projet" # Where projects are created
|
||
DEFAULT_DEPART_DIR = Path.home() / "PycharmProjects" / "structure" / "départ" # Source folder with pole sub-folders
|
||
GEMINI_API_KEY = os.environ.get("AIzaSyDq2LmX_fwKGAwxGAtmBfX940vT2wDQzBU") # Set via env var
|
||
GEMINI_MODEL = "models/gemini-2.5-flash"
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# STEP 1 – PROJECT FOLDER CREATION
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
SUB_FOLDERS = ["compilation", "ingénierie", "RSA-UDS", "plan"]
|
||
|
||
def create_project_structure(project_name: str, base_dir: Path) -> Path:
|
||
"""
|
||
Create:
|
||
<base_dir>/<project_name>/
|
||
compilation/
|
||
ingénierie/
|
||
RSA-UDS/
|
||
plan/
|
||
Returns the project root path.
|
||
"""
|
||
project_root = base_dir / project_name
|
||
project_root.mkdir(parents=True, exist_ok=True)
|
||
|
||
for folder in SUB_FOLDERS:
|
||
(project_root / folder).mkdir(exist_ok=True)
|
||
|
||
log.info(f"✅ Project structure created: {project_root}")
|
||
return project_root
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# STEP 2 – PARSE SOURCE POLE FOLDERS
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
# Pattern: <seq><d|sd><anchors>-<pole_ref> <address>
|
||
# e.g. 1d1-PBO345 rue saint-eugene
|
||
# 2sd2-PBO678 boul. des érables
|
||
FOLDER_PATTERN = re.compile(
|
||
r"^(?P<seq>\d+)"
|
||
r"(?P<dero>sd|d)"
|
||
r"(?P<anchors>\d+)"
|
||
r"-"
|
||
r"(?P<pole_ref>[^\s]+)"
|
||
r"(?:\s+(?P<address>.+))?$",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
def parse_pole_folder_name(folder_name: str) -> dict | None:
|
||
"""
|
||
Parse a source folder name and return a dict with extracted components,
|
||
or None if the name doesn't match the expected pattern.
|
||
"""
|
||
m = FOLDER_PATTERN.match(folder_name)
|
||
if not m:
|
||
return None
|
||
return {
|
||
"seq" : int(m.group("seq")),
|
||
"dero" : m.group("dero").lower(), # 'sd' or 'd'
|
||
"anchors" : int(m.group("anchors")),
|
||
"pole_ref" : m.group("pole_ref"), # e.g. PBO345
|
||
"address" : (m.group("address") or "").strip(),
|
||
"original" : folder_name,
|
||
}
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# STEP 3 – RENAMING & MOVING LOGIC (sd vs d)
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def compute_new_folder_name(parsed: dict, current_seq: int) -> str:
|
||
"""
|
||
Apply business naming rules:
|
||
sd → 10-PBO345 (prefix = current_seq, no address)
|
||
d → 11-PBO345 rue saint-x (prefix = current_seq, with address)
|
||
|
||
NOTE: The leading '1' before 0/1 comes from your spec (10- / 11-).
|
||
If you want the sequence itself used as prefix, adjust here.
|
||
"""
|
||
if parsed["dero"] == "sd":
|
||
return f"{current_seq}0-{parsed['pole_ref']}"
|
||
else:
|
||
addr = f" {parsed['address']}" if parsed["address"] else ""
|
||
return f"{current_seq}1-{parsed['pole_ref']}{addr}"
|
||
|
||
|
||
def rename_excel_files(folder_path: Path, new_folder_name: str):
|
||
"""Rename any .xlsx / .xls inside folder to match new_folder_name."""
|
||
for f in folder_path.iterdir():
|
||
if f.suffix.lower() in (".xlsx", ".xls"):
|
||
new_name = new_folder_name + f.suffix
|
||
f.rename(folder_path / new_name)
|
||
log.info(f" 📄 Excel renamed → {new_name}")
|
||
|
||
|
||
def process_depart_to_rsa_uds(depart_dir: Path, rsa_uds_dir: Path) -> list[Path]:
|
||
"""
|
||
Read all sub-folders in départ, sort by sequence number,
|
||
apply renaming logic, move to RSA-UDS.
|
||
Returns list of destination folder paths (for further processing).
|
||
"""
|
||
entries = []
|
||
for item in depart_dir.iterdir():
|
||
if item.is_dir():
|
||
parsed = parse_pole_folder_name(item.name)
|
||
if parsed:
|
||
entries.append((parsed, item))
|
||
else:
|
||
log.warning(f"⚠️ Unrecognised folder name (skipped): {item.name}")
|
||
|
||
# Sort by original sequence
|
||
entries.sort(key=lambda x: x[0]["seq"])
|
||
|
||
current_seq = 1
|
||
moved_folders = []
|
||
|
||
for parsed, src_path in entries:
|
||
new_name = compute_new_folder_name(parsed, current_seq)
|
||
dst_path = rsa_uds_dir / new_name
|
||
|
||
shutil.copytree(src_path, dst_path)
|
||
log.info(f"📁 {src_path.name!r:40s} → {new_name!r}")
|
||
|
||
# Rename Excel files inside the destination folder
|
||
rename_excel_files(dst_path, new_name)
|
||
|
||
moved_folders.append(dst_path)
|
||
|
||
# Advance counter
|
||
if parsed["dero"] == "sd":
|
||
current_seq += 1
|
||
else:
|
||
current_seq += 2 # Extra jump for CDC document
|
||
|
||
return moved_folders
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# STEP 4 – GEMINI AI IMAGE IDENTIFICATION
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def _encode_image(image_path: Path) -> str:
|
||
"""Return base64-encoded image content."""
|
||
with open(image_path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode("utf-8")
|
||
|
||
|
||
def identify_images_with_gemini(folder_path: Path) -> dict:
|
||
"""
|
||
Send all JPG images in folder to Gemini and ask it to label each as
|
||
'pole' or 'anchor'.
|
||
|
||
Returns:
|
||
{
|
||
"pole": ["IMG_001.jpg", ...],
|
||
"anchor": ["IMG_002.jpg", ...],
|
||
"unknown": [...]
|
||
}
|
||
"""
|
||
jpg_files = sorted(folder_path.glob("*.jpg")) + sorted(folder_path.glob("*.JPG")) + \
|
||
sorted(folder_path.glob("*.jpeg"))
|
||
|
||
if not jpg_files:
|
||
log.warning(f" No JPG images found in {folder_path.name}")
|
||
return {"pole": [], "anchor": [], "unknown": []}
|
||
|
||
if not GEMINI_AVAILABLE or not GEMINI_API_KEY:
|
||
log.warning(" Gemini unavailable – using filename heuristic fallback.")
|
||
return _fallback_image_classification(jpg_files)
|
||
|
||
genai.configure(api_key=GEMINI_API_KEY)
|
||
model = genai.GenerativeModel(GEMINI_MODEL)
|
||
|
||
# Build the content list
|
||
file_names = [f.name for f in jpg_files]
|
||
prompt = (
|
||
"You are analysing images for a utility pole engineering project. "
|
||
"For each image provided, classify it as either 'pole' or 'anchor'. "
|
||
"A pole image shows a utility/telephone pole. "
|
||
"An anchor image shows a ground anchor or guy-wire anchor attachment. "
|
||
f"The filenames are: {file_names}. "
|
||
"Respond ONLY with a JSON object where each key is a filename and "
|
||
"the value is either 'pole' or 'anchor'. Example: "
|
||
'{"IMG_001.jpg": "pole", "IMG_002.jpg": "anchor"}'
|
||
)
|
||
|
||
content = [prompt]
|
||
for img_path in jpg_files:
|
||
content.append({
|
||
"mime_type": "image/jpeg",
|
||
"data": _encode_image(img_path),
|
||
})
|
||
|
||
try:
|
||
response = model.generate_content(content)
|
||
raw = response.text.strip()
|
||
# Strip markdown code fences if present
|
||
raw = re.sub(r"^```json\s*|```$", "", raw, flags=re.MULTILINE).strip()
|
||
import json
|
||
classification = json.loads(raw)
|
||
except Exception as e:
|
||
log.error(f" Gemini error: {e} – falling back to heuristic.")
|
||
return _fallback_image_classification(jpg_files)
|
||
|
||
result = {"pole": [], "anchor": [], "unknown": []}
|
||
for fname, label in classification.items():
|
||
label = label.lower()
|
||
if label in result:
|
||
result[label].append(fname)
|
||
else:
|
||
result["unknown"].append(fname)
|
||
|
||
log.info(f" 🤖 Gemini classified: pole={result['pole']}, anchor={result['anchor']}")
|
||
return result
|
||
|
||
|
||
def _fallback_image_classification(jpg_files: list[Path]) -> dict:
|
||
"""
|
||
Simple heuristic fallback when Gemini is unavailable.
|
||
Files with 'pole' or 'p' in name → pole; 'anchor' or 'a' → anchor.
|
||
Remaining: first half = pole, second half = anchor.
|
||
"""
|
||
result = {"pole": [], "anchor": [], "unknown": []}
|
||
unclassified = []
|
||
|
||
for f in jpg_files:
|
||
lower = f.name.lower()
|
||
if "pole" in lower or lower.startswith("p"):
|
||
result["pole"].append(f.name)
|
||
elif "anchor" in lower or lower.startswith("a"):
|
||
result["anchor"].append(f.name)
|
||
else:
|
||
unclassified.append(f.name)
|
||
|
||
mid = len(unclassified) // 2
|
||
result["pole"].extend(unclassified[:mid])
|
||
result["anchor"].extend(unclassified[mid:])
|
||
|
||
return result
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# STEP 5 – FILE CONVERSION & FINAL SORTING
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def images_to_pdf(image_paths: list[Path], output_pdf: Path):
|
||
"""Combine one or more JPG images into a single PDF page-per-image."""
|
||
if not image_paths:
|
||
return
|
||
|
||
pil_images = []
|
||
for p in image_paths:
|
||
img = Image.open(p).convert("RGB")
|
||
pil_images.append(img)
|
||
|
||
if pil_images:
|
||
first = pil_images[0]
|
||
rest = pil_images[1:]
|
||
first.save(output_pdf, save_all=True, append_images=rest)
|
||
log.info(f" 🖼→📄 {[p.name for p in image_paths]} → {output_pdf.name}")
|
||
|
||
|
||
def excel_to_pdf(excel_path: Path, output_pdf: Path):
|
||
"""
|
||
Convert an Excel file to a PDF using ReportLab.
|
||
Reads all sheets and renders them as tables.
|
||
"""
|
||
wb = load_workbook(excel_path, data_only=True)
|
||
doc = SimpleDocTemplate(
|
||
str(output_pdf),
|
||
pagesize=landscape(letter),
|
||
rightMargin=20, leftMargin=20,
|
||
topMargin=30, bottomMargin=20,
|
||
)
|
||
styles = getSampleStyleSheet()
|
||
story = []
|
||
|
||
for sheet_name in wb.sheetnames:
|
||
ws = wb[sheet_name]
|
||
story.append(Paragraph(f"Sheet: {sheet_name}", styles["Heading2"]))
|
||
story.append(Spacer(1, 6))
|
||
|
||
data = []
|
||
for row in ws.iter_rows(values_only=True):
|
||
data.append([str(cell) if cell is not None else "" for cell in row])
|
||
|
||
if not data:
|
||
continue
|
||
|
||
col_count = max(len(r) for r in data)
|
||
col_width = (landscape(letter)[0] - 40) / max(col_count, 1)
|
||
|
||
table = Table(data, colWidths=[col_width] * col_count, repeatRows=1)
|
||
table.setStyle(TableStyle([
|
||
("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#2C4770")),
|
||
("TEXTCOLOR", (0, 0), (-1, 0), colors.white),
|
||
("FONTSIZE", (0, 0), (-1, -1), 7),
|
||
("GRID", (0, 0), (-1, -1), 0.25, colors.grey),
|
||
("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, colors.HexColor("#F0F4FA")]),
|
||
("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
|
||
]))
|
||
story.append(table)
|
||
story.append(Spacer(1, 20))
|
||
|
||
doc.build(story)
|
||
log.info(f" 📊→📄 {excel_path.name} → {output_pdf.name}")
|
||
|
||
|
||
def process_pole_folder(
|
||
pole_folder: Path,
|
||
compilation_dir: Path,
|
||
pdf_seq: int,
|
||
) -> int:
|
||
"""
|
||
For a single pole folder (already in RSA-UDS):
|
||
1. Identify images via Gemini
|
||
2. Convert pole images → PDF (pdf_seq)
|
||
3. Convert anchor images → PDF (pdf_seq + 1)
|
||
4. Convert Excel → PDF (pdf_seq + 2)
|
||
5. Move all PDFs to compilation
|
||
|
||
Returns the next available pdf_seq.
|
||
"""
|
||
folder_label = pole_folder.name # e.g. "10-PBO345"
|
||
|
||
# Strip leading numeric prefix for base label e.g. "PBO345"
|
||
base_label = re.sub(r"^\d+-", "", folder_label)
|
||
|
||
# ── Identify images ──────────────────────────────────────────────────────
|
||
classification = identify_images_with_gemini(pole_folder)
|
||
|
||
# ── Pole PDF ─────────────────────────────────────────────────────────────
|
||
pole_images = [pole_folder / n for n in classification["pole"] if (pole_folder / n).exists()]
|
||
if pole_images:
|
||
pole_pdf_name = f"{pdf_seq}-{base_label}.pdf"
|
||
images_to_pdf(pole_images, compilation_dir / pole_pdf_name)
|
||
pdf_seq += 1
|
||
|
||
# ── Anchor PDF ───────────────────────────────────────────────────────────
|
||
anchor_images = [pole_folder / n for n in classification["anchor"] if (pole_folder / n).exists()]
|
||
if anchor_images:
|
||
anchor_pdf_name = f"{pdf_seq}-{base_label}.pdf"
|
||
images_to_pdf(anchor_images, compilation_dir / anchor_pdf_name)
|
||
pdf_seq += 1
|
||
|
||
# ── Excel → PDF ──────────────────────────────────────────────────────────
|
||
for xlsx_file in list(pole_folder.glob("*.xlsx")) + list(pole_folder.glob("*.xls")):
|
||
excel_pdf_name = f"{pdf_seq}-{base_label}.pdf"
|
||
excel_to_pdf(xlsx_file, compilation_dir / excel_pdf_name)
|
||
pdf_seq += 1
|
||
|
||
return pdf_seq
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# MAIN ORCHESTRATOR
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def run(project_name: str, base_dir: Path, depart_dir: Path):
|
||
log.info("=" * 65)
|
||
log.info(f" PROJECT : {project_name}")
|
||
log.info(f" BASE DIR: {base_dir}")
|
||
log.info(f" DÉPART : {depart_dir}")
|
||
log.info("=" * 65)
|
||
|
||
# 1 – Create structure
|
||
project_root = create_project_structure(project_name, base_dir)
|
||
rsa_uds_dir = project_root / "RSA-UDS"
|
||
compilation_dir = project_root / "compilation"
|
||
|
||
# 2 & 3 – Move and rename pole folders
|
||
if not depart_dir.exists():
|
||
log.error(f"Source départ folder not found: {depart_dir}")
|
||
sys.exit(1)
|
||
|
||
moved_folders = process_depart_to_rsa_uds(depart_dir, rsa_uds_dir)
|
||
log.info(f"📂 {len(moved_folders)} pole folder(s) moved to RSA-UDS.")
|
||
|
||
# 4 & 5 – Process each pole folder
|
||
pdf_seq = 1
|
||
for folder in sorted(moved_folders):
|
||
log.info(f"\n🔧 Processing: {folder.name}")
|
||
pdf_seq = process_pole_folder(folder, compilation_dir, pdf_seq)
|
||
|
||
log.info("\n" + "=" * 65)
|
||
log.info(f"✅ DONE. All PDFs are in: {compilation_dir}")
|
||
log.info("=" * 65)
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# CLI ENTRY POINT
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Utility pole engineering project automation script",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
Examples:
|
||
python pole_automation.py E65-BK7-26-132679
|
||
python pole_automation.py E65-BK7-26-132679 --base-dir /Users/me/Projects
|
||
python pole_automation.py E65-BK7-26-132679 --depart-dir /Users/me/départ
|
||
python pole_automation.py E65-BK7-26-132679 --gemini-key AIza...
|
||
""",
|
||
)
|
||
parser.add_argument("project_name", help="Project name (e.g. E65-BK7-26-132679)")
|
||
parser.add_argument("--base-dir", type=Path, default=DEFAULT_BASE_DIR,
|
||
help=f"Root directory for projects (default: {DEFAULT_BASE_DIR})")
|
||
parser.add_argument("--depart-dir", type=Path, default=DEFAULT_DEPART_DIR,
|
||
help=f"Source 'départ' folder (default: {DEFAULT_DEPART_DIR})")
|
||
parser.add_argument("--gemini-key", type=str, default=GEMINI_API_KEY,
|
||
help="Gemini API key (or set GEMINI_API_KEY env var)")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Inject Gemini key if provided via CLI
|
||
if args.gemini_key:
|
||
global GEMINI_API_KEY
|
||
GEMINI_API_KEY = args.gemini_key
|
||
|
||
run(args.project_name, args.base_dir, args.depart_dir)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |