File size: 6,181 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 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 132 133 134 135 136 137 138 139 140 141 142 |
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)
# Generate PDF report
pdf_report = PDFReportGenerator(csv_generator.output_csv, self.output_folder, 'damage_estimation_report.pdf')
pdf_report.generate_report()
# # Add download button for the CSV file
# csv_data = df.to_csv(index=False).encode('utf-8')
# st.download_button(
# label="Download CSV",
# data=csv_data,
# file_name='damage_estimation.csv',
# mime='text/csv',
# )
# Add download button for the PDF 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()
|