aryrk commited on
Commit
4c9b339
·
1 Parent(s): fe4b148

[fix error]

Browse files
Files changed (1) hide show
  1. app.py +8 -11
app.py CHANGED
@@ -28,10 +28,11 @@ def reflection_removal(input_image):
28
  if not input_image.lower().endswith((".jpg", ".jpeg", ".png")):
29
  return ["File is not supported (only .jpg, .jpeg, .png)."]
30
 
31
- input_file_name = os.path.basename(input_image)
32
- file_path = os.path.join(UPLOAD_DIR, input_file_name)
33
  shutil.copy(input_image, file_path)
34
 
 
 
35
  cmd = [
36
  "python", "test.py",
37
  "--dataroot", UPLOAD_DIR,
@@ -44,18 +45,14 @@ def reflection_removal(input_image):
44
  ]
45
  subprocess.run(cmd, check=True)
46
 
47
- output_images = []
48
-
49
  for root, _, files in os.walk(RESULTS_DIR):
50
  for file in files:
51
- if file.startswith(input_file_name.split('.')[0]) and file.endswith("_fake.png"):
52
  result_path = os.path.join(root, file)
53
- output_images.append(Image.open(result_path))
54
-
55
- if output_images:
56
- return output_images
57
- else:
58
- return ["No results found."]
59
 
60
  def use_sample_image(sample_image_name):
61
  sample_image_path = os.path.join(SAMPLE_DIR, sample_image_name)
 
28
  if not input_image.lower().endswith((".jpg", ".jpeg", ".png")):
29
  return ["File is not supported (only .jpg, .jpeg, .png)."]
30
 
31
+ file_path = os.path.join(UPLOAD_DIR, os.path.basename(input_image))
 
32
  shutil.copy(input_image, file_path)
33
 
34
+ input_filename = os.path.splitext(os.path.basename(file_path))[0]
35
+
36
  cmd = [
37
  "python", "test.py",
38
  "--dataroot", UPLOAD_DIR,
 
45
  ]
46
  subprocess.run(cmd, check=True)
47
 
 
 
48
  for root, _, files in os.walk(RESULTS_DIR):
49
  for file in files:
50
+ if file.endswith("_fake.png") and file.startswith(input_filename):
51
  result_path = os.path.join(root, file)
52
+ return [Image.open(result_path)]
53
+
54
+ return ["No results found."]
55
+
 
 
56
 
57
  def use_sample_image(sample_image_name):
58
  sample_image_path = os.path.join(SAMPLE_DIR, sample_image_name)