|
import streamlit as st
|
|
from PIL import Image
|
|
import os
|
|
import shutil
|
|
from ultralytics import YOLO
|
|
import tempfile
|
|
import pandas as pd
|
|
import plotly.express as px
|
|
from fetch_original import FileProcessor
|
|
from Damage_calculation import DamageCalculator
|
|
from generator import CSVGenerator
|
|
from pdf_report import PDFReportGenerator
|
|
|
|
class SegmentationModel:
|
|
def __init__(self, model_path):
|
|
self.model = YOLO(model_path)
|
|
|
|
def predict(self, input_path):
|
|
result = self.model.predict(source=input_path, save=True, show_conf=False, conf=0.70, save_txt=True)
|
|
img_path = result[0].save_dir
|
|
return img_path
|
|
|
|
class SegmentationApp:
|
|
def __init__(self, car_parts_model_path, damage_model_path, output_folder):
|
|
self.car_parts_model = SegmentationModel(car_parts_model_path)
|
|
self.damage_model = SegmentationModel(damage_model_path)
|
|
self.output_folder = output_folder
|
|
|
|
def copy_folder(self, src_folder, dest_folder):
|
|
if not os.path.exists(dest_folder):
|
|
os.makedirs(dest_folder)
|
|
for item in os.listdir(src_folder):
|
|
s = os.path.join(src_folder, item)
|
|
d = os.path.join(dest_folder, item)
|
|
if os.path.isdir(s):
|
|
shutil.copytree(s, d, dirs_exist_ok=True)
|
|
else:
|
|
shutil.copy2(s, d)
|
|
|
|
def clean_output_folder(self):
|
|
parts_output_folder = os.path.join(self.output_folder, 'parts')
|
|
damage_output_folder = os.path.join(self.output_folder, 'damage')
|
|
if os.path.exists(parts_output_folder):
|
|
shutil.rmtree(parts_output_folder)
|
|
if os.path.exists(damage_output_folder):
|
|
shutil.rmtree(damage_output_folder)
|
|
|
|
def run(self):
|
|
st.title("Car Damage Analyses And Cost Estimation")
|
|
st.markdown("### Important Guidelines for Using the App")
|
|
st.write("""
|
|
<div style="background-color: pink; color: black; padding: 10px; border-radius: 5px;">
|
|
<ul>
|
|
<li>Provide 4 images of the car from the following views: front, back, left, and right.</li>
|
|
<li>Ensure the images are taken at proper angles, standing parallel to the car, and that they are clear.</li>
|
|
<li>Make sure there are no obstacles between the car and the camera.</li>
|
|
<li>The car should not be behind any objects in the images.</li>
|
|
</ul>
|
|
</div>
|
|
""", unsafe_allow_html=True)
|
|
|
|
uploaded_files = st.file_uploader("Choose up to 4 images...", type=["jpg", "jpeg", "png", "webp"], accept_multiple_files=True)
|
|
|
|
if uploaded_files and len(uploaded_files) == 4:
|
|
st.write("Running segmentation models on uploaded images...")
|
|
car_parts_results_paths = []
|
|
damage_results_paths = []
|
|
|
|
self.clean_output_folder()
|
|
|
|
with st.spinner("Analyzing total damage..."):
|
|
for i, uploaded_file in enumerate(uploaded_files):
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
|
|
temp_file.write(uploaded_file.getvalue())
|
|
temp_image_path = temp_file.name
|
|
|
|
car_parts_results = self.car_parts_model.predict(temp_image_path)
|
|
car_parts_results_paths.append(car_parts_results)
|
|
|
|
damage_results = self.damage_model.predict(temp_image_path)
|
|
damage_results_paths.append(damage_results)
|
|
|
|
self.car_parts_results_paths = car_parts_results_paths
|
|
self.damage_results_paths = damage_results_paths
|
|
|
|
for car_parts_results in car_parts_results_paths:
|
|
parts_output_folder = os.path.join(self.output_folder, 'parts')
|
|
self.copy_folder(car_parts_results, parts_output_folder)
|
|
|
|
for damage_results in damage_results_paths:
|
|
damage_output_folder = os.path.join(self.output_folder, 'damage')
|
|
self.copy_folder(damage_results, damage_output_folder)
|
|
|
|
file_processor = FileProcessor()
|
|
file_processor.process_output_folder(self.output_folder)
|
|
|
|
csv_generator = CSVGenerator(self.output_folder)
|
|
damage_data = csv_generator.process_files()
|
|
|
|
df = pd.read_csv(csv_generator.output_csv)
|
|
st.write("Damage Estimation Report")
|
|
st.dataframe(df)
|
|
|
|
|
|
pdf_report = PDFReportGenerator(csv_generator.output_csv, self.output_folder, 'damage_estimation_report.pdf')
|
|
pdf_report.generate_report()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
with open(pdf_report.pdf_path, 'rb') as f:
|
|
pdf_data = f.read()
|
|
st.download_button(
|
|
label="Download PDF Report",
|
|
data=pdf_data,
|
|
file_name='damage_estimation_report.pdf',
|
|
mime='application/pdf',
|
|
)
|
|
|
|
damage_df = pd.DataFrame(damage_data)
|
|
fig = px.pie(damage_df, values='coverage_percentage', names='damage_class', title='Total Damage Percentages by Type')
|
|
st.plotly_chart(fig)
|
|
|
|
st.success("Analysis Completed.")
|
|
|
|
elif uploaded_files:
|
|
st.error("Please upload exactly four images.")
|
|
|
|
if __name__ == "__main__":
|
|
CAR_PARTS_MODEL_PATH = 'model/car_parts.pt'
|
|
DAMAGE_MODEL_PATH = 'model/damage_4_class.pt'
|
|
OUTPUT_FOLDER = 'Output'
|
|
app = SegmentationApp(CAR_PARTS_MODEL_PATH, DAMAGE_MODEL_PATH, OUTPUT_FOLDER)
|
|
app.run()
|
|
|
