Alex Strick van Linschoten commited on
Commit
513a0c5
1 Parent(s): 411008f

update tmp_dir

Browse files
Files changed (1) hide show
  1. app.py +28 -7
app.py CHANGED
@@ -33,6 +33,7 @@ labels = learn.dls.vocab
33
 
34
 
35
  def predict(pdf, confidence, generate_file):
 
36
  document = fitz.open(pdf.name)
37
  results = []
38
  images = []
@@ -47,12 +48,18 @@ def predict(pdf, confidence, generate_file):
47
  if probs[0] > (confidence / 100):
48
  redaction_count = len(images)
49
  image_pixmap.save(
50
- os.path.join(tmp_dir, pdf.name, f"page-{page_num}.png")
 
 
51
  )
52
  images.append(
53
  [
54
  f"Redacted page #{redaction_count + 1} on page {page_num}",
55
- os.path.join(tmp_dir, pdf.name, f"page-{page_num}.png"),
 
 
 
 
56
  ]
57
  )
58
 
@@ -61,12 +68,18 @@ def predict(pdf, confidence, generate_file):
61
  for page in range(len(results))
62
  if results[page]["redacted"] > (confidence / 100)
63
  ]
64
- report = os.path.join(tmp_dir, pdf.name, "redacted_pages.pdf")
 
 
65
  if generate_file:
66
  pdf = FPDF()
67
  pdf.set_auto_page_break(0)
68
  imagelist = sorted(
69
- [i for i in os.listdir(tmp_dir, pdf.name) if i.endswith("png")]
 
 
 
 
70
  )
71
  for image in imagelist:
72
  # with PILImage.open(os.path.join(tmp_dir, image)) as img:
@@ -76,7 +89,9 @@ def predict(pdf, confidence, generate_file):
76
  # else:
77
  # pdf.add_page("P")
78
  # pdf.image(os.path.join(tmp_dir, image))
79
- with PILImage.open(os.path.join(tmp_dir, pdf.name, image)) as img:
 
 
80
  size = img.size
81
  if size[0] > size[1]:
82
  pdf.add_page("L")
@@ -95,9 +110,15 @@ def predict(pdf, confidence, generate_file):
95
  label_color="#FF59D6",
96
  )
97
  pred_dict["img"].save(
98
- os.path.join(tmp_dir, pdf.name, f"pred-{image}")
 
 
 
 
 
 
99
  )
100
- pdf.image(os.path.join(tmp_dir, pdf.name, f"pred-{image}"))
101
  pdf.output(report, "F")
102
 
103
  text_output = f"A total of {len(redacted_pages)} pages were redacted. \n\n The redacted page numbers were: {', '.join(redacted_pages)}."
33
 
34
 
35
  def predict(pdf, confidence, generate_file):
36
+ filename_without_extension = pdf.name[:-4]
37
  document = fitz.open(pdf.name)
38
  results = []
39
  images = []
48
  if probs[0] > (confidence / 100):
49
  redaction_count = len(images)
50
  image_pixmap.save(
51
+ os.path.join(
52
+ tmp_dir, filename_without_extension, f"page-{page_num}.png"
53
+ )
54
  )
55
  images.append(
56
  [
57
  f"Redacted page #{redaction_count + 1} on page {page_num}",
58
+ os.path.join(
59
+ tmp_dir,
60
+ filename_without_extension,
61
+ f"page-{page_num}.png",
62
+ ),
63
  ]
64
  )
65
 
68
  for page in range(len(results))
69
  if results[page]["redacted"] > (confidence / 100)
70
  ]
71
+ report = os.path.join(
72
+ tmp_dir, filename_without_extension, "redacted_pages.pdf"
73
+ )
74
  if generate_file:
75
  pdf = FPDF()
76
  pdf.set_auto_page_break(0)
77
  imagelist = sorted(
78
+ [
79
+ i
80
+ for i in os.listdir(tmp_dir, filename_without_extension)
81
+ if i.endswith("png")
82
+ ]
83
  )
84
  for image in imagelist:
85
  # with PILImage.open(os.path.join(tmp_dir, image)) as img:
89
  # else:
90
  # pdf.add_page("P")
91
  # pdf.image(os.path.join(tmp_dir, image))
92
+ with PILImage.open(
93
+ os.path.join(tmp_dir, filename_without_extension, image)
94
+ ) as img:
95
  size = img.size
96
  if size[0] > size[1]:
97
  pdf.add_page("L")
110
  label_color="#FF59D6",
111
  )
112
  pred_dict["img"].save(
113
+ os.path.join(
114
+ tmp_dir, filename_without_extension, f"pred-{image}"
115
+ )
116
+ )
117
+ pdf.image(
118
+ os.path.join(
119
+ tmp_dir, filename_without_extension, f"pred-{image}"
120
  )
121
+ )
122
  pdf.output(report, "F")
123
 
124
  text_output = f"A total of {len(redacted_pages)} pages were redacted. \n\n The redacted page numbers were: {', '.join(redacted_pages)}."