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') os.system('pip install pycocotools @ git+https://github.com/philferriere/cocoapi.git@2929bd2ef6b451054755dfd7ceb09278f935f7ad#subdirectory=PythonAPI') 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 = [(0,0,0,1), (0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1), (0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1),(0,0,0,1), (0,0,0,1), (0,0,0,1), (0,0,0,1),(255,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('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)) seg = seg.convert("RGBA") pixdata = seg.load() width, height = seg.size for y in range(height): for x in range(width): if pixdata[x, y] == (255, 255, 255, 255): pixdata[x, y] = (255, 255, 255, 0) return 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()