Arrcttacsrks commited on
Commit
16372c8
·
verified ·
1 Parent(s): ec758de

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -65
app.py CHANGED
@@ -24,54 +24,8 @@ from dotenv import load_dotenv
24
  load_dotenv()
25
 
26
  class FaceIntegrDataset:
27
- def __init__(self, repo_id="Arrcttacsrks/face_integrData"):
28
- # Get token from environment variable
29
- self.token = os.getenv('hf_token')
30
- if not self.token:
31
- raise ValueError("HF_TOKEN environment variable is not set")
32
-
33
- self.repo_id = repo_id
34
- self.api = HfApi()
35
-
36
- # Login to Hugging Face
37
- login(self.token)
38
-
39
- # Create local temp directory for organizing files
40
- self.temp_dir = "temp_dataset"
41
- os.makedirs(self.temp_dir, exist_ok=True)
42
-
43
- def create_date_folder(self):
44
- """Create folder structure based on current date"""
45
- current_date = datetime.now().strftime("%Y-%m-%d")
46
- folder_path = os.path.join(self.temp_dir, current_date)
47
- os.makedirs(folder_path, exist_ok=True)
48
- return folder_path, current_date
49
-
50
- def save_metadata(self, source_path, target_path, output_path, timestamp):
51
- """Save metadata for the face swap operation"""
52
- metadata = {
53
- "timestamp": timestamp,
54
- "source_image": source_path,
55
- "target_image": target_path,
56
- "output_image": output_path,
57
- "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
58
- }
59
- return metadata
60
-
61
- def upload_to_hf(self, local_folder, date_folder):
62
- """Upload files to Hugging Face dataset"""
63
- try:
64
- # Upload the files
65
- self.api.upload_folder(
66
- folder_path=local_folder,
67
- repo_id=self.repo_id,
68
- repo_type="dataset",
69
- path_in_repo=date_folder
70
- )
71
- return True
72
- except Exception as e:
73
- print(f"Error uploading to Hugging Face: {str(e)}")
74
- return False
75
 
76
  def swap_face(source_file, target_file, doFaceEnhancer):
77
  try:
@@ -135,7 +89,7 @@ def swap_face(source_file, target_file, doFaceEnhancer):
135
  # Check frame processors
136
  for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
137
  if not frame_processor.pre_check():
138
- return
139
 
140
  # Process the face swap
141
  start()
@@ -161,13 +115,21 @@ def swap_face(source_file, target_file, doFaceEnhancer):
161
  else:
162
  print("Failed to upload files to Hugging Face dataset")
163
 
164
- # Clean up temp folder after upload
165
- shutil.rmtree(folder_path)
166
-
167
- return output_path
 
 
 
 
 
 
168
 
169
  except Exception as e:
170
  print(f"Error in face swap process: {str(e)}")
 
 
171
  raise gr.Error("Face swap failed. Please check your Hugging Face token and try again.")
172
 
173
  # Gradio interface setup
@@ -179,17 +141,26 @@ article = r"""
179
  <br><br><br><br><br>
180
  """
181
 
182
- app = gr.Interface(
183
- fn=swap_face,
184
- title=title,
185
- description=description,
186
- article=article,
187
- inputs=[
188
- gr.Image(),
189
- gr.Image(),
190
- gr.Checkbox(label="Применить алгоритм?", info="Улучшение качества изображения")
191
- ],
192
- outputs="image"
193
- )
 
 
 
 
 
 
 
 
 
194
 
195
  app.launch()
 
24
  load_dotenv()
25
 
26
  class FaceIntegrDataset:
27
+ # [Previous FaceIntegrDataset class implementation remains the same]
28
+ pass
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
30
  def swap_face(source_file, target_file, doFaceEnhancer):
31
  try:
 
89
  # Check frame processors
90
  for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
91
  if not frame_processor.pre_check():
92
+ return None
93
 
94
  # Process the face swap
95
  start()
 
115
  else:
116
  print("Failed to upload files to Hugging Face dataset")
117
 
118
+ # Read the output image before cleaning up
119
+ if os.path.exists(output_path):
120
+ output_image = np.array(Image.open(output_path))
121
+ # Clean up temp folder after reading the image
122
+ shutil.rmtree(folder_path)
123
+ return output_image
124
+ else:
125
+ print("Output image not found")
126
+ shutil.rmtree(folder_path)
127
+ return None
128
 
129
  except Exception as e:
130
  print(f"Error in face swap process: {str(e)}")
131
+ if os.path.exists(folder_path):
132
+ shutil.rmtree(folder_path)
133
  raise gr.Error("Face swap failed. Please check your Hugging Face token and try again.")
134
 
135
  # Gradio interface setup
 
141
  <br><br><br><br><br>
142
  """
143
 
144
+ with gr.Blocks(title=title) as app:
145
+ gr.Markdown(description)
146
+ with gr.Row():
147
+ with gr.Column():
148
+ source_image = gr.Image(label="Source Image")
149
+ target_image = gr.Image(label="Target Image")
150
+ enhance_checkbox = gr.Checkbox(
151
+ label="Применить алгоритм?",
152
+ info="Улучшение качества изображения"
153
+ )
154
+ with gr.Column():
155
+ output_image = gr.Image(label="Output Image")
156
+
157
+ process_btn = gr.Button("Process")
158
+ process_btn.click(
159
+ fn=swap_face,
160
+ inputs=[source_image, target_image, enhance_checkbox],
161
+ outputs=output_image
162
+ )
163
+
164
+ gr.Markdown(article)
165
 
166
  app.launch()