File size: 2,704 Bytes
dba26ff |
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 |
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)
|