Alex Strick van Linschoten commited on
Commit
411008f
1 Parent(s): 36149a3

update tmp_dir

Browse files
Files changed (1) hide show
  1. app.py +12 -9
app.py CHANGED
@@ -36,8 +36,7 @@ def predict(pdf, confidence, generate_file):
36
  document = fitz.open(pdf.name)
37
  results = []
38
  images = []
39
- tmp_dir = tempfile.TemporaryDirectory().name
40
- # tmp_dir = tempfile.gettempdir()
41
  for page_num, page in enumerate(document, start=1):
42
  image_pixmap = page.get_pixmap()
43
  image = image_pixmap.tobytes()
@@ -47,11 +46,13 @@ def predict(pdf, confidence, generate_file):
47
  )
48
  if probs[0] > (confidence / 100):
49
  redaction_count = len(images)
50
- image_pixmap.save(os.path.join(tmp_dir, f"page-{page_num}.png"))
 
 
51
  images.append(
52
  [
53
  f"Redacted page #{redaction_count + 1} on page {page_num}",
54
- os.path.join(tmp_dir, f"page-{page_num}.png"),
55
  ]
56
  )
57
 
@@ -60,12 +61,12 @@ def predict(pdf, confidence, generate_file):
60
  for page in range(len(results))
61
  if results[page]["redacted"] > (confidence / 100)
62
  ]
63
- report = os.path.join(tmp_dir, "redacted_pages.pdf")
64
  if generate_file:
65
  pdf = FPDF()
66
  pdf.set_auto_page_break(0)
67
  imagelist = sorted(
68
- [i for i in os.listdir(tmp_dir) if i.endswith("png")]
69
  )
70
  for image in imagelist:
71
  # with PILImage.open(os.path.join(tmp_dir, image)) as img:
@@ -75,7 +76,7 @@ def predict(pdf, confidence, generate_file):
75
  # else:
76
  # pdf.add_page("P")
77
  # pdf.image(os.path.join(tmp_dir, image))
78
- with PILImage.open(os.path.join(tmp_dir, image)) as img:
79
  size = img.size
80
  if size[0] > size[1]:
81
  pdf.add_page("L")
@@ -93,8 +94,10 @@ def predict(pdf, confidence, generate_file):
93
  font_size=16,
94
  label_color="#FF59D6",
95
  )
96
- pred_dict["img"].save(os.path.join(tmp_dir, f"pred-{image}"))
97
- pdf.image(os.path.join(tmp_dir, f"pred-{image}"))
 
 
98
  pdf.output(report, "F")
99
 
100
  text_output = f"A total of {len(redacted_pages)} pages were redacted. \n\n The redacted page numbers were: {', '.join(redacted_pages)}."
 
36
  document = fitz.open(pdf.name)
37
  results = []
38
  images = []
39
+ tmp_dir = tempfile.gettempdir()
 
40
  for page_num, page in enumerate(document, start=1):
41
  image_pixmap = page.get_pixmap()
42
  image = image_pixmap.tobytes()
 
46
  )
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
  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
  # 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")
 
94
  font_size=16,
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)}."