cdnuts commited on
Commit
02d04ca
·
verified ·
1 Parent(s): 75afda4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +82 -17
app.py CHANGED
@@ -176,8 +176,7 @@ class ImageDataset(Dataset):
176
  img_path = self.image_files[idx]
177
  img = Image.open(img_path).convert('RGB')
178
  return self.transform(img), os.path.basename(img_path)
179
-
180
-
181
  @spaces.GPU(duration=299)
182
  def process_images(images, threshold):
183
  dataset = ImageDataset(images, transform)
@@ -221,24 +220,75 @@ def process_zip(zip_file, threshold):
221
  if zip_file is None:
222
  return None, None
223
 
224
- temp_dir = "temp_images"
225
- os.makedirs(temp_dir, exist_ok=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
226
 
227
- with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
228
- zip_ref.extractall(temp_dir)
 
 
 
 
 
229
 
230
- all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
231
- image_files = [f for f in all_files if is_valid_image(f)]
232
- results = process_images(image_files, threshold)
233
-
234
- temp_file = NamedTemporaryFile(delete=False, suffix=".zip")
235
- with zipfile.ZipFile(temp_file, "w") as zip_ref:
236
- for image_name, text_no_impl, _ in results:
237
- with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:
238
- file.write(text_no_impl.encode())
239
- temp_file.seek(0)
240
- df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
241
 
 
 
 
 
 
 
 
 
242
  return temp_file.name, df
243
 
244
  with gr.Blocks(css=".output-class { display: none; }") as demo:
@@ -286,6 +336,21 @@ with gr.Blocks(css=".output-class { display: none; }") as demo:
286
  inputs=[zip_input, multi_threshold_slider],
287
  outputs=[zip_output, dataframe_output]
288
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
289
 
290
  if __name__ == "__main__":
291
  demo.queue().launch()
 
176
  img_path = self.image_files[idx]
177
  img = Image.open(img_path).convert('RGB')
178
  return self.transform(img), os.path.basename(img_path)
179
+
 
180
  @spaces.GPU(duration=299)
181
  def process_images(images, threshold):
182
  dataset = ImageDataset(images, transform)
 
220
  if zip_file is None:
221
  return None, None
222
 
223
+ with tempfile.TemporaryDirectory() as temp_dir:
224
+ with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
225
+ zip_ref.extractall(temp_dir)
226
+
227
+ all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
228
+ image_files = [f for f in all_files if is_valid_image(f)]
229
+ results = process_images(image_files, threshold)
230
+
231
+ temp_file = NamedTemporaryFile(delete=False, suffix=".zip")
232
+ with zipfile.ZipFile(temp_file, "w") as zip_ref:
233
+ for image_name, text_no_impl, _ in results:
234
+ with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:
235
+ file.write(text_no_impl.encode())
236
+ temp_file.seek(0)
237
+ df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])
238
 
239
+ return temp_file.name, df
240
+
241
+ @spaces.GPU(duration=120) # Reduced GPU duration for less wait time...
242
+ def process_images_light(images, threshold):
243
+ dataset = ImageDataset(images, transform)
244
+
245
+ dataloader = DataLoader(dataset, batch_size=16, num_workers=0, pin_memory=True, drop_last=False)
246
 
247
+ all_results = []
248
+
249
+ with torch.no_grad():
250
+ for batch, filenames in dataloader:
251
+
252
+ batch = batch.to(device)
253
+ with torch.no_grad():
254
+ logits = model(batch)
255
+ probabilities = torch.nn.functional.sigmoid(logits)
256
+
257
+ for i, prob in enumerate(probabilities):
258
+ indices = torch.where(prob > threshold)[0]
259
+ values = prob[indices]
260
+
261
+ temp = []
262
+ tag_score = dict()
263
+ for j in range(indices.size(0)):
264
+ temp.append([allowed_tags[indices[j]], values[j].item()])
265
+ tag_score[allowed_tags[indices[j]]] = values[j].item()
266
+
267
+ tags = ", ".join([t[0] for t in temp])
268
+ all_results.append((filenames[i], tags, tag_score))
269
+
270
+ return all_results
271
+
272
+ def process_zip_light(zip_file, threshold):
273
+ if zip_file is None:
274
+ return None, None
275
+
276
+ with tempfile.TemporaryDirectory() as temp_dir:
277
+ with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
278
+ zip_ref.extractall(temp_dir)
279
+
280
+ all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
281
+ image_files = [f for f in all_files if is_valid_image(f)]
282
+ results = process_images_light(image_files, threshold)
283
 
284
+ temp_file = NamedTemporaryFile(delete=False, suffix=".zip")
285
+ with zipfile.ZipFile(temp_file, "w") as zip_ref:
286
+ for image_name, text_no_impl, _ in results:
287
+ with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:
288
+ file.write(text_no_impl.encode())
289
+ temp_file.seek(0)
290
+ df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])
291
+
292
  return temp_file.name, df
293
 
294
  with gr.Blocks(css=".output-class { display: none; }") as demo:
 
336
  inputs=[zip_input, multi_threshold_slider],
337
  outputs=[zip_output, dataframe_output]
338
  )
339
+ with gr.TabItem("Multiple Images (Light)"):
340
+ with gr.Row():
341
+ with gr.Column():
342
+ zip_input_light = gr.File(label="Upload ZIP file", file_types=['.zip'])
343
+ multi_threshold_slider_light = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
344
+ process_button_light = gr.Button("Process Images (Light)")
345
+ with gr.Column():
346
+ zip_output_light = gr.File(label="Download Tagged Text Files (ZIP)")
347
+ dataframe_output_light = gr.Dataframe(label="Image Tags Summary")
348
+
349
+ process_button_light.click(
350
+ fn=process_zip_light,
351
+ inputs=[zip_input_light, multi_threshold_slider_light],
352
+ outputs=[zip_output_light, dataframe_output_light]
353
+ )
354
 
355
  if __name__ == "__main__":
356
  demo.queue().launch()