aryrk commited on
Commit
f6fd486
·
1 Parent(s): d3aec7f

[feat] delete specific result after being processed

Browse files
Files changed (1) hide show
  1. app.py +13 -3
app.py CHANGED
@@ -51,18 +51,28 @@ def reflection_removal(input_image, preprocess_type="resize_and_crop", session_i
51
  ]
52
  subprocess.run(cmd, check=True)
53
 
 
54
  for root, _, files in os.walk(RESULTS_DIR):
55
  for file in files:
56
  if file.startswith(input_filename) and file.endswith("_fake.png"):
57
  result_path = os.path.join(root, file)
58
  output_image = Image.open(result_path)
59
-
60
- clear_session_files(session_id)
61
- return output_image
 
 
 
 
 
62
 
63
  clear_session_files(session_id)
 
 
 
64
  return "No results found."
65
 
 
66
  def clear_session_files(session_id):
67
  session_dir = os.path.join("./sessions", session_id)
68
  if os.path.exists(session_dir):
 
51
  ]
52
  subprocess.run(cmd, check=True)
53
 
54
+ result_files = []
55
  for root, _, files in os.walk(RESULTS_DIR):
56
  for file in files:
57
  if file.startswith(input_filename) and file.endswith("_fake.png"):
58
  result_path = os.path.join(root, file)
59
  output_image = Image.open(result_path)
60
+ result_files.append(result_path)
61
+
62
+ os.remove(result_path)
63
+ elif file.startswith(input_filename) and file.endswith("_real.png"):
64
+ real_path = os.path.join(root, file)
65
+ result_files.append(real_path)
66
+
67
+ os.remove(real_path)
68
 
69
  clear_session_files(session_id)
70
+
71
+ if result_files:
72
+ return output_image
73
  return "No results found."
74
 
75
+
76
  def clear_session_files(session_id):
77
  session_dir = os.path.join("./sessions", session_id)
78
  if os.path.exists(session_dir):