Update app.py

app.py

@@ -1,7 +1,7 @@
 try:
     import detectron2
 except:
-    import os 
+    import os
     os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 
 from matplotlib.pyplot import axis
@@ -10,6 +10,7 @@ import requests
 import numpy as np
 from torch import nn
 import requests
+import cv2
 
 import torch
 import detectron2
@@ -19,6 +20,8 @@ from detectron2.config import get_cfg
 from detectron2.utils.visualizer import Visualizer
 from detectron2.data import MetadataCatalog
 from detectron2.utils.visualizer import ColorMode
+from detectron2.structures import Instances
+from detectron2.structures import Boxes
 
 damage_model_path = 'damage/model_final.pth'
 scratch_model_path = 'scratch/model_final.pth'
@@ -28,7 +31,7 @@ if torch.cuda.is_available():
     device = 'cuda'
 else:
     device = 'cpu'
-
+ 
 cfg_scratches = get_cfg()
 cfg_scratches.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
 cfg_scratches.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
@@ -83,6 +86,7 @@ metadata_parts.thing_classes = ['_background_',
  'trunk',
  'wheel']
 
+
 def merge_segment(pred_segm):
     merge_dict = {}
     for i in range(len(pred_segm)):
@@ -90,36 +94,35 @@ def merge_segment(pred_segm):
         for j in range(i+1,len(pred_segm)):
             if torch.sum(pred_segm[i]*pred_segm[j])>0:
                 merge_dict[i].append(j)
-    
+
     to_delete = []
     for key in merge_dict:
         for element in merge_dict[key]:
             to_delete.append(element)
-    
+
     for element in to_delete:
         merge_dict.pop(element,None)
-        
+
     empty_delete = []
     for key in merge_dict:
         if merge_dict[key] == []:
             empty_delete.append(key)
-    
+
     for element in empty_delete:
         merge_dict.pop(element,None)
-        
+
     for key in merge_dict:
         for element in merge_dict[key]:
             pred_segm[key]+=pred_segm[element]
-            
+
     except_elem = list(set(to_delete))
-    
+
     new_indexes = list(range(len(pred_segm)))
     for elem in except_elem:
         new_indexes.remove(elem)
-        
+
     return pred_segm[new_indexes]
 
-    
 def inference(image):
     img = np.array(image)
     outputs_damage = predictor_damage(img)
@@ -136,7 +139,8 @@ def inference(image):
     parts_classes = parts_data['pred_classes']
     new_inst = detectron2.structures.Instances((1024,1024))
     new_inst.set('pred_masks',merge_segment(out_dict['pred_masks']))
-    
+
+
     parts_damage_dict = {}
     parts_list_damages = []
     for part in parts_classes:
@@ -145,7 +149,7 @@ def inference(image):
         for i in range(len(parts_masks)):
             if torch.sum(parts_masks[i]*mask)>0:
                 parts_damage_dict[metadata_parts.thing_classes[parts_classes[i]]].append('scratch')
-                parts_list_damages.append(f'{metadata_parts.thing_classes[parts_classes[i]]} has scratch')              
+                parts_list_damages.append(f'{metadata_parts.thing_classes[parts_classes[i]]} has scratch')
                 print(f'{metadata_parts.thing_classes[parts_classes[i]]} has scratch')
     for mask in merged_damage_masks:
         for i in range(len(parts_masks)):
@@ -154,35 +158,86 @@ def inference(image):
                 parts_list_damages.append(f'{metadata_parts.thing_classes[parts_classes[i]]} has damage')
                 print(f'{metadata_parts.thing_classes[parts_classes[i]]} has damage')
 
+    # Define the colors for the scratch and damage masks
+    scratch_color = (0, 0, 255)  # red
+    damage_color = (0, 255, 255)  # yellow
+    # Convert the scratch and damage masks to numpy arrays
+    scratch_masks_arr = np.array(scratch_masks)
+    damage_masks_arr = np.array(damage_masks)
+    # Resize the scratch and damage masks to match the size of the original image
+    scratch_mask_resized = cv2.resize(scratch_masks_arr[0].astype(np.uint8), (img.shape[1], img.shape[0]))
+    damage_mask_resized = cv2.resize(damage_masks_arr[0].astype(np.uint8), (img.shape[1], img.shape[0]))
+    # Merge the scratch and damage masks into a single binary mask
+    merged_mask = np.zeros_like(scratch_mask_resized)
+    merged_mask[(scratch_mask_resized> 0) | (damage_mask_resized > 0)] = 255
+    # Overlay the merged mask on top of the original image
+    overlay = img.copy()
+    overlay[merged_mask == 255] = (0, 255, 0)  # green color for the merged mask
+    overlay[damage_mask_resized == 255] = damage_color  # yellow color for the damage mask
+    #output = cv2.addWeighted(overlay, 0.5, img, 0.5, 0)
+
+    # Merge the instance predictions from both predictors
+    image_np = np.array(image)
+    height, width, channels = image_np.shape
+     
+    # Get the predicted boxes from the scratches predictor
+    pred_boxes_scratch = outputs_scratch["instances"].pred_boxes.tensor
+
+    # Get the predicted boxes from the damage predictor
+    pred_boxes_damage = outputs_damage["instances"].pred_boxes.tensor
+
+    # Concatenate the predicted boxes along the batch dimension
+    merged_boxes = torch.cat([pred_boxes_scratch, pred_boxes_damage], dim=0)
+
+    # Create a new Instances object with the merged boxes
+    merged_instances = Instances((image_np.shape[0], image_np.shape[1]))
+    merged_instances.pred_boxes = Boxes(merged_boxes)
+  
+
+    # Visualize the Masks
     v_d = Visualizer(img[:, :, ::-1],
-                   metadata=metadata_damage, 
-                   scale=0.5, 
+                   metadata=metadata_damage,
+                   scale=0.5,
                    instance_mode=ColorMode.SEGMENTATION   # remove the colors of unsegmented pixels. This option is only available for segmentation models
     )
-    #v_d = Visualizer(img,scale=1.2)
-    #print(outputs["instances"].to('cpu'))
+    v_d = Visualizer(img,scale=1.2)
     out_d = v_d.draw_instance_predictions(new_inst)
     img1 = out_d.get_image()[:, :, ::-1]
 
     v_s = Visualizer(img[:, :, ::-1],
-                   metadata=metadata_scratch, 
-                   scale=0.5, 
+                   metadata=metadata_scratch,
+                   scale=0.5,
                    instance_mode=ColorMode.SEGMENTATION   # remove the colors of unsegmented pixels. This option is only available for segmentation models
     )
-    #v_s = Visualizer(img,scale=1.2)
+    v_s = Visualizer(img,scale=1.2)
     out_s = v_s.draw_instance_predictions(outputs_scratch["instances"])
     img2 = out_s.get_image()[:, :, ::-1]
 
     v_p = Visualizer(img[:, :, ::-1],
-                   metadata=metadata_parts, 
-                   scale=0.5, 
+                   metadata=metadata_parts,
+                   scale=0.5,
                    instance_mode=ColorMode.SEGMENTATION   # remove the colors of unsegmented pixels. This option is only available for segmentation models
     )
-    #v_p = Visualizer(img,scale=1.2)
+    v_p = Visualizer(img,scale=1.2)
     out_p = v_p.draw_instance_predictions(outputs_parts["instances"])
     img3 = out_p.get_image()[:, :, ::-1]
-    
-    return img1, img2, img3, parts_list_damages
+
+    # Visualize the overlay
+    v_m = Visualizer(overlay[:, :, ::-1],
+                 metadata=metadata_damage,
+                 scale=1.2,
+                 instance_mode=ColorMode.SEGMENTATION   # display the overlay in black and white
+                )
+    # Draw the overlay with instance predictions
+    overlay_with_predictions = v_m.draw_instance_predictions(merged_instances.to("cpu")).get_image()[:, :, ::-1]
+    #v_m = Visualizer(overlay,scale=1.2)
+    out = v_m.draw_instance_predictions(merged_instances)
+    output = out.get_image()[:, :, ::-1]
+
+
+
+
+    return img1, img2, img3, parts_list_damages, output
 
 
 with gr.Blocks() as demo:
@@ -201,13 +256,15 @@ with gr.Blocks() as demo:
                 im3 = gr.Image(type='numpy',label='Image of car parts')
             with gr.Tab('Information about damaged parts'):
                 intersections = gr.Textbox(label='Information about type of damages on each part')
-    
+            with gr.Tab('Image of overlayed damage parts'):
+                overlayed = gr.Image(type='numpy',label='Image of overlayed damage parts')
+
     #actions
     submit_button.click(
         fn=inference,
         inputs = [image],
-        outputs = [im1,im2,im3,intersections]
+        outputs = [im1,im2,im3,intersections, overlayed]
     )
-        
+
 if __name__ == "__main__":
-    demo.launch()
+   demo.launch()