Update app.py

app.py

@@ -22,12 +22,13 @@ def detect(image):
     people_mask = transform(~people_mask)
     people_mask = people_mask.resize((image.width, image.height), resample=Image.Resampling.BILINEAR)
 
-    vehicle_classes = [2, 3, 5, 7] 
+    vehicle_classes = [2, 3, 5, 7] # Classes: car (2), motorcycle (3), bus (5) and truck (7)
     license_plates = list()
-
+    vehicle_boxes = []
     for seg_box in seg_boxes:
         if seg_box[5] in vehicle_classes:
             vehicle_box = seg_box[:4].to(torch.int32)
+            vehicle_boxes.append(vehicle_box)
             vehicle_crop = image.crop(vehicle_box.tolist())
             imgsz = (vehicle_crop.height, vehicle_crop.width) if vehicle_crop.width < 640 and vehicle_crop.height < 640 else (640, 640)
             lp_result = lp_model(vehicle_crop, imgsz=imgsz, device="cpu")[0]
@@ -37,22 +38,31 @@ def detect(image):
                 license_plates.append(torch.add(lp_box, vehicle_offset))            
 
     lp_mask = Image.new(mode="L", size=image.size, color=255)
-    draw = ImageDraw.Draw(lp_mask)
+    lp_draw = ImageDraw.Draw(lp_mask)
 
     for license_plate in license_plates:
-        draw.rectangle(license_plate.tolist(), fill = 0)
+        lp_draw.rectangle(license_plate.tolist(), fill = 0)
+
+    vehicle_mask = Image.new(mode="L", size=image.size, color=255)
+    vehicle_draw = ImageDraw.Draw(vehicle_mask)
+    for vehicle_box in vehicle_boxes:
+        vehicle_draw.rectangle(vehicle_box.tolist(), outline = 0, width=5)
+
+
 
+    #TODO: move combination to caller function
     combined_mask = Image.fromarray(np.minimum.reduce([np.array(m) for m in [people_mask, lp_mask]]))
-    return combined_mask
+    return combined_mask, people_mask, lp_mask, vehicle_mask
 
 
 def test_comb(image):
-    mask = detect(image)
+    mask, people_mask, lp_mask, vm = detect(image)
     blurred = image.filter(ImageFilter.GaussianBlur(30))
     anonymized = Image.composite(image, blurred, mask)
     ## TODO: Tempfile statt einem generischen File
     anonymized.save("anon.JPG")
-    return "anon.JPG"
+    annotation_list = [(1 - np.asarray(people_mask) / 255, "Person"), (1 - np.asarray(vm) / 255, "Fahrzeug"), (1 - np.asarray(lp_mask) / 255, "Kennzeichen")]
+    return "anon.JPG", (image, annotation_list)
 
 
 css = """
@@ -77,8 +87,8 @@ Ein Prototyp des it@M InnovationLab (itm.innolab@muenchen.de)
 demo_upload = gr.Interface(
     title="Image Anonymizer",
     fn=test_comb, 
-    inputs=gr.Image(type="pil"), 
-    outputs=gr.Image(),
+    inputs=gr.Image(type="pil", label="Zu anonymisierendes Bild"), 
+    outputs=[gr.Image(label="Anonymisiertes Bild"), gr.AnnotatedImage(label="Erkannte Regionen")],
     allow_flagging="never",
     examples="examples",
     description=description,