Create new file

app.py

@@ -0,0 +1,183 @@
+import segmentation_models as sm
+import numpy as np
+import os
+import cv2
+import keras
+import matplotlib.colors as colorsHTML
+from PIL import Image
+import gradio as gr
+
+import os
+os.system('wget https://huggingface.co/Armandoliv/cars-parts-segmentation-unet-resnet18/resolve/main/best_model.h5')
+
+c= ['_background_', 'back_bumper', 'back_glass', 'back_left_door','back_left_light',
+ 'back_right_door', 'back_right_light', 'front_bumper','front_glass',  
+ 'front_left_door',  'front_left_light', 'front_right_door',  'front_right_light',  'hood',  'left_mirror',
+ 'right_mirror', 'tailgate',  'trunk', 'wheel']
+
+colors = [ 	(245,255,250), (75,0,130),  (0,255,0), (32,178,170),(0,0,255), (0,255,255), (255,0,255), (128,0,128), (255,140,0),
+          (85,107,47), (102,205,170), (0,191,255), (255,0,0), 	(255,228,196), (205,133,63),
+          (220,20,60), (255,69,0), 	(143,188,143), (255,255,0)]
+
+
+sm.set_framework('tf.keras')
+
+sm.framework()
+
+BACKBONE = 'resnet18'
+n_classes = 19
+activation =  'softmax'
+
+#create model
+model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)
+
+# load best weights
+model.load_weights('/content/best_model.h5') 
+
+def get_colored_segmentation_image(seg_arr, n_classes, colors=colors):
+    output_height = seg_arr.shape[0]
+    output_width = seg_arr.shape[1]
+
+    seg_img = np.zeros((output_height, output_width, 3))
+
+    for c in range(n_classes):
+        seg_arr_c = seg_arr[:, :] == c
+          # print(sum(sum(seg_arr_c)), colors[c] )
+        seg_img[:, :, 0] += ((seg_arr_c)*(colors[c][0])).astype('uint8')
+        seg_img[:, :, 1] += ((seg_arr_c)*(colors[c][1])).astype('uint8')
+        seg_img[:, :, 2] += ((seg_arr_c)*(colors[c][2])).astype('uint8')
+
+    return seg_img/255
+
+def get_legends(class_names, colors, tags):
+
+    n_classes = len(class_names)
+    legend = np.zeros(((len(class_names) * 25) + 25, 125, 3),
+                      dtype="uint8") + 255
+
+    class_names_colors = enumerate(zip(class_names[:n_classes],
+                                       colors[:n_classes]))
+    j = 0
+    for (i, (class_name, color)) in class_names_colors:
+        if i in tags:
+          color = [int(c) for c in color]
+          cv2.putText(legend, class_name, (5, (j * 25) + 17),
+                      cv2.FONT_HERSHEY_COMPLEX, 0.35, (0, 0, 0), 1)
+          cv2.rectangle(legend, (100, (j* 25)), (125, (j * 25) + 25),
+                        tuple(color), -1)
+          j +=1
+    return legend
+
+
+
+def preprocess_image(path_img):
+  img = Image.open(path_img)
+  ww = 512
+  hh = 512
+  img.thumbnail((hh, ww))
+  i = np.array(img)
+  ht, wd, cc= i.shape
+
+  # create new image of desired size and color (blue) for padding
+  color = (0,0,0)
+  result = np.full((hh,ww,cc), color, dtype=np.uint8)
+
+  # copy img image into center of result image
+  result[:ht, :wd] = img
+  return result, ht, wd
+
+def concat_lengends(seg_img, legend_img):
+
+  new_h = np.maximum(seg_img.shape[0], legend_img.shape[0])
+  new_w = seg_img.shape[1] + legend_img.shape[1]
+
+  out_img = np.zeros((new_h, new_w, 3)).astype('uint8') + legend_img[0, 0, 0]
+
+  out_img[:legend_img.shape[0], :  legend_img.shape[1]] = np.copy(legend_img)
+  out_img[:seg_img.shape[0], legend_img.shape[1]:] = np.copy(seg_img)
+
+  return out_img
+
+def main_convert(filename):
+
+  print(filename)
+  #load the image
+  img_path = filename
+  img = Image.open(img_path).convert("RGB")
+  tags = []
+
+  #preprocess the image
+  img_scaled_arr = preprocess_image(img_path)
+  image = np.expand_dims(img_scaled_arr[0], axis=0)
+
+  #make the predictions
+  pr_mask = model.predict(image).squeeze()
+  pr_mask_int = np.zeros((pr_mask.shape[0],pr_mask.shape[1]))
+
+  #filter the smallest noisy segments
+  kernel = np.ones((5, 5), 'uint8')
+
+  for i in range(1,19):
+    array_one = np.round(pr_mask[:,:,i])
+    op = cv2.morphologyEx(array_one, cv2.MORPH_OPEN, kernel)
+    if sum(sum(op ==1)) > 100:
+      tags.append(i)
+      pr_mask_int[op ==1] = i 
+
+  img_segmented = np.array(Image.fromarray(pr_mask_int[:img_scaled_arr[1], :img_scaled_arr[2]]).resize(img.size))
+
+  seg = get_colored_segmentation_image(img_segmented,19, colors=colors)
+
+  fused_img = ((np.array(img)/255)/2 + seg/2).astype('float32')
+
+  seg = Image.fromarray((seg*255).astype(np.uint8))
+  fused_img  = Image.fromarray((fused_img *255).astype(np.uint8))
+
+  #get the legends
+  legend_predicted = get_legends(c, colors, tags)
+
+  final_img = concat_lengends(np.array(fused_img), np.array(legend_predicted))
+
+  return final_img, seg
+  
+  
+  
+inputs = [gr.Image(type="filepath", label="Car Image")]
+outputs = [gr.Image(type="PIL.Image", label="Detected Segments Image"),gr.Image(type="PIL.Image", label="Segment Image")]
+
+
+title = "Car Parts Segmentation APP"
+description = """This demo uses AI Models to detect 18 parts of cars: \n
+1: background, 
+2: back bumper, 
+3: back glass, 
+4: back left door, 
+5: back left light, 
+6: back right door, 
+7: back right light, 
+8: front bumper, 
+9: front glass, 
+10: front left door, 
+11: front left light, 
+12: front right door, 
+13: front right light, 
+14: hood, 
+15: left mirror, 
+16: right mirror, 
+17: tailgate, 
+18: trunk, 
+19: wheel"""
+
+examples = [['test_image.jpeg']]
+io = gr.Interface(fn=main_convert, inputs=inputs, outputs=outputs, title=title, description=description, examples=examples,
+                  css= """.gr-button-primary { background: -webkit-linear-gradient( 
+                    90deg, #355764 0%, #55a8a1 100% ) !important;     background: #355764;
+                        background: linear-gradient( 
+                    90deg, #355764 0%, #55a8a1 100% ) !important;
+                        background: -moz-linear-gradient( 90deg, #355764 0%, #55a8a1 100% ) !important;
+                        background: -webkit-linear-gradient( 
+                    90deg, #355764 0%, #55a8a1 100% ) !important;
+                    color:white !important}"""
+                  )
+                  
+io.launch()