Update app.py

app.py

@@ -9,67 +9,81 @@ transform = torchvision.transforms.ToPILImage()
 seg_model = YOLO("yolov8m-seg.pt")
 lp_model = YOLO("yolov8m_lp.pt")
 
-def detect_person(image: Image):
-    result = seg_model(image, device="CPU")[0]
-    masks = result.masks.data
-    clss = result.boxes.cls
-    person_indices = torch.where(clss == 0)
-    person_masks = masks[person_indices]
+
+def detect(image):
+    seg_result = seg_model(image, device="CPU")[0]
+    seg_masks = seg_result.masks.data
+    seg_clss = seg_result.boxes.cls
+    seg_boxes = seg_result.boxes.data
+
+    person_indices = torch.where(seg_clss == 0)
+    person_masks = seg_masks[person_indices]
     people_mask = torch.any(person_masks, dim=0).to(torch.uint8) * 255
-    mask = transform(~people_mask)
-    mask = mask.resize((image.width, image.height), resample=Image.Resampling.BILINEAR)
-    return mask
-
-def detect_license_plate(image: Image):
-    result = lp_model(image, imgsz=(image.height, image.width), device="cpu")[0]
-    boxes = result.boxes.data[:, :4]
-    mask = Image.new(mode="L", size=image.size, color=255)
-    draw = ImageDraw.Draw(mask)
-    for box in boxes:
-        draw.rectangle(list(box), fill=0)
-    return mask
-
-
-def detect_dummy(image: Image):
-    return Image.new(mode="L", size=image.size, color=255)
-
-
-detectors = {
-    "Person": detect_person,
-    "License Plate": detect_license_plate
-}
-
-def test_gradio(image):
-    masks = [detect_person(image), detect_license_plate(image)]
-    combined = np.minimum.reduce([np.array(m) for m in masks])
-    mask = Image.fromarray(combined)
-    # Apply blur through mask
+    people_mask = transform(~people_mask)
+    people_mask = people_mask.resize((image.width, image.height), resample=Image.Resampling.BILINEAR)
+
+    vehicle_classes = [2, 3, 5, 7] 
+    license_plates = list()
+
+    for seg_box in seg_boxes:
+        if seg_box[5] in vehicle_classes:
+            vehicle_box = seg_box[:4].to(torch.int32)
+            vehicle_crop = image.crop(vehicle_box.tolist())
+            lp_result = lp_model(vehicle_crop, imgsz=(vehicle_crop.height, vehicle_crop.width), device="cpu")[0]
+            lp_boxes = lp_result.boxes.data[:, :4]
+            vehicle_offset = torch.cat((vehicle_box[:2], vehicle_box[:2]))
+            for lp_box in lp_boxes:
+                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)
+
+    for license_plate in license_plates:
+        draw.rectangle(license_plate.tolist(), fill = 0)
+
+    combined_mask = Image.fromarray(np.minimum.reduce([np.array(m) for m in [people_mask, lp_mask]]))
+    return combined_mask
+
+
+def test_comb(image):
+    mask = 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"
 
-# demo_live = gr.Interface(
-#     fn=test_gradio, 
-#     inputs=gr.Image(source="webcam", type="pil", shape=(640, 480)), 
-#     outputs=gr.Image(type="pil")
-# )
+
+css = """
+P { text-align: center }
+H3 { text-align: center }
+"""
+
+description = """
+### ML-Prototyp zur Anonymisierung von Bildern
+Es werden Personen sowie Kennzeichen zensiert.
+Große Bilder können einige Zeit benötigen.
+"""
+
+article = """
+Nutzt YOLOv8-Modelle zur Erkennung / Segmentierung der Bilder.
+
+Code: https://huggingface.co/spaces/it-at-m/image-anonymizer/tree/main
+
+Ein Prototyp des it@M InnovationLab (itm.innolab@muenchen.de)
+"""
 
 demo_upload = gr.Interface(
-    fn=test_gradio, 
+    title="Image Anonymizer",
+    fn=test_comb, 
     inputs=gr.Image(type="pil"), 
-    outputs=gr.Image()
+    outputs=gr.Image(),
+    allow_flagging="never",
+    examples="examples",
+    description=description,
+    article=article,
+    css=css
 )
 
-# demo = gr.TabbedInterface(
-#     interface_list=[demo_live, demo_upload], 
-#     tab_names=["Webcam", "Bild hochladen"], 
-#     title="Image Anonymizer"
-# )
-
-# print(__name__)
-# demo_upload.launch(server_name="localhost", server_port=8080)
-# demo.launch(server_name="localhost", server_port=8080)
-
+demo_upload.queue(concurrency_count=1)
 demo_upload.launch()