add preds

app.py

@@ -1,14 +1,30 @@
+import os
+import tempfile
+
+import fitz
 import gradio as gr
+import PIL
 import skimage
 from fastai.learner import load_learner
 from fastai.vision.all import *
-from huggingface_hub import hf_hub_download
-import fitz
-import tempfile
-import os
 from fpdf import FPDF
+from huggingface_hub import hf_hub_download
+from icevision.all import *
+from icevision.models.checkpoint import *
 from PIL import Image as PILImage
 
+checkpoint_path = "./2022-01-15-vfnet-post-self-train.pth"
+checkpoint_and_model = model_from_checkpoint(checkpoint_path)
+model = checkpoint_and_model["model"]
+model_type = checkpoint_and_model["model_type"]
+class_map = checkpoint_and_model["class_map"]
+
+img_size = checkpoint_and_model["img_size"]
+valid_tfms = tfms.A.Adapter(
+    [*tfms.A.resize_and_pad(img_size), tfms.A.Normalize()]
+)
+
+
 learn = load_learner(
     hf_hub_download("strickvl/redaction-classifier-fastai", "model.pkl")
 )
@@ -51,16 +67,37 @@ def predict(pdf, confidence, generate_file):
             [i for i in os.listdir(tmp_dir) if i.endswith("png")]
         )
         for image in imagelist:
+            # with PILImage.open(os.path.join(tmp_dir, image)) as img:
+            #     size = img.size
+            #     if size[0] > size[1]:
+            #         pdf.add_page("L")
+            #     else:
+            #         pdf.add_page("P")
+            # pdf.image(os.path.join(tmp_dir, image))
             with PILImage.open(os.path.join(tmp_dir, image)) as img:
                 size = img.size
                 if size[0] > size[1]:
                     pdf.add_page("L")
                 else:
                     pdf.add_page("P")
-            # pdf.image(os.path.join(tmp_dir, image), w=190, h=280)
-            pdf.image(os.path.join(tmp_dir, image))
+                pred_dict = model_type.end2end_detect(
+                    img,
+                    valid_tfms,
+                    model,
+                    class_map=class_map,
+                    detection_threshold=0.7,
+                    display_label=True,
+                    display_bbox=True,
+                    return_img=True,
+                    font_size=16,
+                    label_color="#FF59D6",
+                )
+                pred_dict["img"].save(os.path.join(tmp_dir, f"pred-{image}"))
+            pdf.image(os.path.join(tmp_dir, f"pred-{image}"))
         pdf.output(report, "F")
+
     text_output = f"A total of {len(redacted_pages)} pages were redacted. \n\n The redacted page numbers were: {', '.join(redacted_pages)}."
+
     if generate_file:
         return text_output, images, report
     else: