Mohaddz commited on
Commit
d311746
·
verified ·
1 Parent(s): ba38e71

app and requirements

Browse files
Files changed (2) hide show
  1. app.py +83 -0
  2. requirements.txt +0 -0
app.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from PIL import Image
4
+ import numpy as np
5
+ import tensorflow as tf
6
+ from transformers import SegformerForSemanticSegmentation, AutoFeatureExtractor
7
+ import cv2
8
+ import json
9
+
10
+ # Load models
11
+ part_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/huggingCars")
12
+ damage_seg_model = SegformerForSemanticSegmentation.from_pretrained("Mohaddz/DamageSeg")
13
+ feature_extractor = AutoFeatureExtractor.from_pretrained("Mohaddz/huggingCars")
14
+ dl_model = tf.keras.models.load_model('improved_car_damage_prediction_model.h5')
15
+
16
+ # Load parts list
17
+ with open('cars117.json', 'r', encoding='utf-8') as f:
18
+ data = json.load(f)
19
+ all_parts = sorted(list(set(part for entry in data.values() for part in entry.get('replaced_parts', []))))
20
+
21
+ def process_image(image):
22
+ # Convert to RGB if it's not
23
+ if image.mode != 'RGB':
24
+ image = image.convert('RGB')
25
+
26
+ # Prepare input for the model
27
+ inputs = feature_extractor(images=image, return_tensors="pt")
28
+
29
+ # Get damage segmentation
30
+ with torch.no_grad():
31
+ damage_output = damage_seg_model(**inputs).logits
32
+ damage_features = damage_output.squeeze().detach().numpy()
33
+
34
+ # Create damage segmentation heatmap
35
+ damage_heatmap = create_heatmap(damage_features)
36
+ damage_heatmap_resized = cv2.resize(damage_heatmap, (image.size[0], image.size[1]))
37
+
38
+ # Create annotated damage image
39
+ image_array = np.array(image)
40
+ damage_mask = np.argmax(damage_features, axis=0)
41
+ damage_mask_resized = cv2.resize(damage_mask, (image.size[0], image.size[1]), interpolation=cv2.INTER_NEAREST)
42
+ overlay = np.zeros_like(image_array)
43
+ overlay[damage_mask_resized > 0] = [255, 0, 0] # Red color for damage
44
+ annotated_image = cv2.addWeighted(image_array, 1, overlay, 0.5, 0)
45
+
46
+ # Process for part prediction and heatmap
47
+ with torch.no_grad():
48
+ part_output = part_seg_model(**inputs).logits
49
+ part_features = part_output.squeeze().detach().numpy()
50
+ part_heatmap = create_heatmap(part_features)
51
+ part_heatmap_resized = cv2.resize(part_heatmap, (image.size[0], image.size[1]))
52
+
53
+ # Predict parts to replace
54
+ input_vector = np.concatenate([part_features.mean(axis=(1, 2)), damage_features.mean(axis=(1, 2))])
55
+ prediction = dl_model.predict(np.array([input_vector]))
56
+ predicted_parts = [(all_parts[i], float(prob)) for i, prob in enumerate(prediction[0]) if prob > 0.1]
57
+ predicted_parts.sort(key=lambda x: x[1], reverse=True)
58
+
59
+ return (Image.fromarray(annotated_image),
60
+ Image.fromarray(damage_heatmap_resized),
61
+ Image.fromarray(part_heatmap_resized),
62
+ "\n".join([f"{part}: {prob:.2f}" for part, prob in predicted_parts[:5]]))
63
+
64
+ def create_heatmap(features):
65
+ heatmap = np.sum(features, axis=0)
66
+ heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())
67
+ heatmap = np.uint8(255 * heatmap)
68
+ return cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
69
+
70
+ iface = gr.Interface(
71
+ fn=process_image,
72
+ inputs=gr.Image(type="pil"),
73
+ outputs=[
74
+ gr.Image(type="pil", label="Annotated Damage"),
75
+ gr.Image(type="pil", label="Damage Heatmap"),
76
+ gr.Image(type="pil", label="Part Segmentation Heatmap"),
77
+ gr.Textbox(label="Predicted Parts to Replace")
78
+ ],
79
+ title="Car Damage Assessment",
80
+ description="Upload an image of a damaged car to get an assessment."
81
+ )
82
+
83
+ iface.launch(share=True)
requirements.txt ADDED
Binary file (4.15 kB). View file