|
import os
|
|
from PIL import Image
|
|
|
|
class FileProcessor:
|
|
PARTS_NAMES = [
|
|
"Quarter-panel", "Front-wheel", "Back-window", "Trunk", "Front-door", "Rocker-panel",
|
|
"Grille", "Windshield", "Front-window", "Back-door", "Headlight", "Back-wheel",
|
|
"Back-windshield", "Hood", "Fender", "Tail-light", "License-plate", "Front-bumper",
|
|
"Back-bumper", "Mirror", "Roof"
|
|
]
|
|
|
|
DAMAGE_NAMES = [
|
|
"Broken part", "Scratch", "Dent", "Paint chip"
|
|
]
|
|
|
|
@staticmethod
|
|
def convert_normalized_to_original(coords, width, height):
|
|
original_coords = []
|
|
for i in range(0, len(coords), 2):
|
|
x = coords[i] * width
|
|
y = coords[i + 1] * height
|
|
original_coords.extend([x, y])
|
|
return original_coords
|
|
|
|
def process_file(self, file_path, class_names, image_size):
|
|
width, height = image_size
|
|
with open(file_path, "r") as f:
|
|
lines = f.readlines()
|
|
|
|
new_lines = []
|
|
for line in lines:
|
|
parts = line.strip().split()
|
|
class_index = int(parts[0])
|
|
coords = list(map(float, parts[1:]))
|
|
original_coords = self.convert_normalized_to_original(coords, width, height)
|
|
class_name = class_names[class_index]
|
|
new_line = f"{class_name} " + " ".join(map(str, original_coords)) + "\n"
|
|
new_lines.append(new_line)
|
|
|
|
with open(file_path, "w") as f:
|
|
f.writelines(new_lines)
|
|
|
|
def process_folder(self, folder_path, class_names):
|
|
labels_folder = os.path.join(folder_path, "labels")
|
|
for file in os.listdir(folder_path):
|
|
if file.endswith(".jpg"):
|
|
image_path = os.path.join(folder_path, file)
|
|
if os.path.exists(image_path):
|
|
image = Image.open(image_path)
|
|
image_size = image.size
|
|
txt_file = file.replace(".jpg", ".txt")
|
|
txt_file_path = os.path.join(labels_folder, txt_file)
|
|
if os.path.exists(txt_file_path):
|
|
self.process_file(txt_file_path, class_names, image_size)
|
|
else:
|
|
print(f"Text file {txt_file_path} not found.")
|
|
else:
|
|
print(f"Image {image_path} not found.")
|
|
|
|
def process_output_folder(self, output_folder):
|
|
parts_folder = os.path.join(output_folder, "parts")
|
|
damage_folder = os.path.join(output_folder, "damage")
|
|
|
|
self.process_folder(parts_folder, self.PARTS_NAMES)
|
|
self.process_folder(damage_folder, self.DAMAGE_NAMES)
|
|
|
