AnonymousECCV15285/MMIB_Counterfactual_image_generation_tool
0
1import streamlit as st
2import os
3import sys
4import tempfile
5import zipfile
6import json
7import random
8import math
9import csv
10from pathlib import Path
11from datetime import datetime
12import time
13
14sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
15
16script_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'scripts')
17sys.path.insert(0, script_dir)
18
19try:
20 from pipeline import (
21 generate_counterfactuals,
22 generate_base_scene,
23 save_scene,
24 render_scene,
25 create_patched_render_script,
26 IMAGE_COUNTERFACTUALS,
27 NEGATIVE_COUNTERFACTUALS
28 )
29 try:
30 import sys
31 script_dir = os.path.dirname(os.path.abspath(__file__))
32 scripts_path = os.path.join(script_dir, 'scripts')
33 if scripts_path not in sys.path:
34 sys.path.insert(0, scripts_path)
35 from generate_questions_mapping import (
36 load_scene,
37 generate_question_for_scene as _generate_question_for_scene_file,
38 answer_question_for_scene,
39 generate_mapping_with_questions
40 )
41 except ImportError:
42 def load_scene(scene_file):
43 with open(scene_file, 'r') as f:
44 return json.load(f)
45 def answer_question_for_scene(question, scene):
46 objects = scene.get('objects', [])
47 return len(objects)
48 _generate_question_for_scene_file = None
49 generate_mapping_with_questions = None
50 PIPELINE_AVAILABLE = True
51except ImportError as e:
52 print(f"Warning: Error importing pipeline functions: {e}")
53 PIPELINE_AVAILABLE = False
54 answer_question_for_scene = None
55 load_scene = None
56 _generate_question_for_scene_file = None
57
58st.set_page_config(
59 page_title="Counterfactual Image Generator",
60 page_icon="🎨",
61 layout="wide",
62 initial_sidebar_state="expanded"
63)
64
65st.markdown("""
66 <style>
67 .main-header {
68 font-size: 2.5rem;
69 font-weight: bold;
70 color: #1f77b4;
71 text-align: center;
72 margin-bottom: 2rem;
73 }
74 .stButton>button {
75 width: 100%;
76 height: 3.5rem;
77 font-size: 1.2rem;
78 font-weight: bold;
79 background-color: #1f77b4;
80 color: white;
81 border-radius: 0.5rem;
82 }
83 .stButton>button:hover {
84 background-color: #1565c0;
85 }
86 .info-box {
87 padding: 1rem;
88 border-radius: 0.5rem;
89 background-color: #f0f2f6;
90 margin: 1rem 0;
91 }
92 </style>
93""", unsafe_allow_html=True)
94
95def create_zip_file(output_dir, zip_path):
96 with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
97 for root, dirs, files in os.walk(output_dir):
98 for file in files:
99 file_path = os.path.join(root, file)
100 arcname = os.path.relpath(file_path, output_dir)
101 zipf.write(file_path, arcname)
102
103
104def generate_fallback_scene(num_objects, scene_idx):
105 script_dir = os.path.dirname(os.path.abspath(__file__))
106 props_path = os.path.join(script_dir, 'data', 'properties.json')
107
108 try:
109 with open(props_path, 'r') as f:
110 properties = json.load(f)
111 except:
112 properties = {
113 'shapes': {'cube': 'SmoothCube_v2', 'sphere': 'Sphere', 'cylinder': 'SmoothCylinder'},
114 'colors': {'gray': [87, 87, 87], 'red': [173, 35, 35], 'blue': [42, 75, 215],
115 'green': [29, 105, 20], 'brown': [129, 74, 25], 'purple': [129, 38, 192],
116 'cyan': [41, 208, 208], 'yellow': [255, 238, 51]},
117 'materials': {'rubber': 'Rubber', 'metal': 'MyMetal'},
118 'sizes': {'large': 0.7, 'small': 0.35}
119 }
120
121 shapes = list(properties['shapes'].keys())
122 colors = list(properties['colors'].keys())
123 materials = list(properties['materials'].keys())
124 sizes = list(properties['sizes'].keys())
125
126 scene_num = scene_idx + 1
127 scene = {
128 'split': 'fallback',
129 'image_index': scene_num,
130 'image_filename': f'scene_{scene_num:04d}_original.png',
131 'objects': [],
132 'directions': {
133 'behind': (0.0, -1.0, 0.0),
134 'front': (0.0, 1.0, 0.0),
135 'left': (-1.0, 0.0, 0.0),
136 'right': (1.0, 0.0, 0.0),
137 'above': (0.0, 0.0, 1.0),
138 'below': (0.0, 0.0, -1.0)
139 }
140 }
141
142 positions = []
143 min_dist = 0.25
144
145 for i in range(num_objects):
146 max_attempts = 100
147 placed = False
148
149 for attempt in range(max_attempts):
150 x = random.uniform(-3, 3)
151 y = random.uniform(-3, 3)
152 z = random.uniform(0.35, 0.7)
153
154 collision = False
155 size = random.choice(sizes)
156 r = properties['sizes'][size]
157
158 for (px, py, pz, pr) in positions:
159 dist = math.sqrt((x - px)**2 + (y - py)**2)
160 if dist < (r + pr + min_dist):
161 collision = True
162 break
163
164 if not collision:
165 positions.append((x, y, z, r))
166 placed = True
167 break
168
169 if not placed:
170 x = random.uniform(-3, 3)
171 y = random.uniform(-3, 3)
172 z = random.uniform(0.35, 0.7)
173 size = random.choice(sizes)
174 r = properties['sizes'][size]
175 positions.append((x, y, z, r))
176
177 shape = random.choice(shapes)
178 color = random.choice(colors)
179 material = random.choice(materials)
180
181 obj = {
182 'shape': shape,
183 'size': size,
184 'material': material,
185 '3d_coords': [x, y, z],
186 'rotation': random.uniform(0, 360),
187 'pixel_coords': [0, 0, 0],
188 'color': color
189 }
190
191 scene['objects'].append(obj)
192
193 return scene
194
195
196def generate_question_for_scene_dict(scene):
197 if _generate_question_for_scene_file is None:
198 objects = scene.get('objects', [])
199 if len(objects) == 0:
200 return "How many objects are in the scene?", {}
201
202 colors = list(set(obj.get('color') for obj in objects if obj.get('color')))
203 shapes = list(set(obj.get('shape') for obj in objects if obj.get('shape')))
204
205 if colors:
206 return f"How many {random.choice(colors)} objects are there?", {'color': random.choice(colors)}
207 else:
208 return "How many objects are in the scene?", {}
209
210 with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tmp_file:
211 json.dump(scene, tmp_file)
212 tmp_path = tmp_file.name
213
214 try:
215 question, params = _generate_question_for_scene_file(tmp_path)
216 return question, params
217 finally:
218 try:
219 os.unlink(tmp_path)
220 except:
221 pass
222
223
224def generate_counterfactual_scenes(num_scenes, num_objects, min_objects, max_objects, num_counterfactuals,
225 cf_types, same_cf_type, min_change_score, max_cf_attempts, min_noise_level,
226 output_dir, blender_path=None, use_gpu=0, samples=512,
227 width=320, height=240, skip_render=False, generate_questions=False,
228 semantic_only=False, negative_only=False):
229 if not PIPELINE_AVAILABLE:
230 return {
231 'success': False,
232 'error': 'Pipeline functions not available. Please ensure pipeline.py is accessible.'
233 }
234
235 scenes_dir = os.path.join(output_dir, 'scenes')
236 images_dir = os.path.join(output_dir, 'images')
237 os.makedirs(scenes_dir, exist_ok=True)
238 os.makedirs(images_dir, exist_ok=True)
239
240 script_dir = os.path.dirname(os.path.abspath(__file__))
241 cwd = os.getcwd()
242 import shutil
243 import time
244
245 temp_output_dir = os.path.join(cwd, 'temp_output')
246 if os.path.exists(temp_output_dir):
247 for attempt in range(3):
248 try:
249 shutil.rmtree(temp_output_dir)
250 break
251 except Exception as e:
252 if attempt < 2:
253 time.sleep(0.3)
254 else:
255 print(f"Warning: Could not remove temp_output after 3 attempts: {e}")
256
257 render_patched_path = os.path.join(cwd, 'render_images_patched.py')
258 if os.path.exists(render_patched_path):
259 for attempt in range(3):
260 try:
261 time.sleep(0.2)
262 if os.path.exists(render_patched_path):
263 os.remove(render_patched_path)
264 break
265 except Exception as e:
266 if attempt < 2:
267 time.sleep(0.3)
268 else:
269 print(f"Warning: Could not remove render_images_patched.py after 3 attempts: {e}")
270
271 blender_available = False
272 if blender_path is None:
273 try:
274 from pipeline import find_blender
275 blender_path = find_blender()
276 except:
277 blender_path = 'blender'
278
279 if blender_path and blender_path != 'blender':
280 blender_available = os.path.exists(blender_path)
281 else:
282 try:
283 import subprocess
284 test_path = blender_path if blender_path and blender_path != 'blender' else 'blender'
285 env = os.environ.copy()
286 result = subprocess.run([test_path, '--version'], capture_output=True, timeout=5, env=env)
287 blender_available = (result.returncode == 0)
288 except:
289 blender_available = False
290
291 successful_scenes = 0
292 successful_renders = 0
293 error_messages = []
294
295 try:
296 for scene_idx in range(num_scenes):
297 if num_objects is not None:
298 scene_num_objects = num_objects
299 else:
300 scene_num_objects = random.randint(min_objects, max_objects)
301
302 base_scene = None
303
304 if blender_available:
305 scene_error = None
306 for retry in range(3):
307 try:
308 import io
309 import contextlib
310 output_buffer = io.StringIO()
311 with contextlib.redirect_stdout(output_buffer), contextlib.redirect_stderr(output_buffer):
312 base_scene = generate_base_scene(
313 scene_num_objects,
314 blender_path,
315 scene_idx
316 )
317 blender_output = output_buffer.getvalue()
318 if blender_output and retry == 2:
319 st.text(f"Blender output for scene {scene_idx + 1} (last 1000 chars):")
320 st.code(blender_output[-1000:] if len(blender_output) > 1000 else blender_output)
321
322 if base_scene and len(base_scene.get('objects', [])) > 0:
323 break
324 elif base_scene is None:
325 if retry == 2:
326 scene_error = f"generate_base_scene returned None - Blender may have failed (check output above)"
327 error_messages.append(f"Scene {scene_idx + 1}: {scene_error}")
328 elif len(base_scene.get('objects', [])) == 0:
329 if retry == 2:
330 scene_error = f"Scene has 0 objects - Blender may have hit max_retries (check output above)"
331 error_messages.append(f"Scene {scene_idx + 1}: {scene_error}")
332 except FileNotFoundError as e:
333 scene_error = f"Blender not found: {e}"
334 error_messages.append(f"Scene {scene_idx + 1}: {scene_error}")
335 blender_available = False
336 break
337 except Exception as e:
338 import traceback
339 scene_error = f"Error generating base scene: {str(e)}"
340 print(f"Error generating base scene (retry {retry + 1}/3): {e}")
341 print(f" Traceback: {traceback.format_exc()}")
342 if retry == 2:
343 full_error = f"Scene {scene_idx + 1}: {scene_error} (Blender path: {blender_path})"
344 error_messages.append(full_error)
345 blender_available = False
346 continue
347 else:
348 print(f"Scene {scene_idx + 1} (Blender not available)...")
349 base_scene = generate_fallback_scene(scene_num_objects, scene_idx)
350
351 if not base_scene or len(base_scene.get('objects', [])) == 0:
352 error_detail = f"Scene {scene_idx + 1}: Failed to generate"
353 if blender_available:
354 error_detail += f" (Blender was available at {blender_path} but returned empty scene)"
355 else:
356 error_detail += " (Blender not available, fallback scene also failed)"
357 print(f"Failed to generate scene {scene_idx + 1}")
358 print(f" Blender available: {blender_available}")
359 print(f" Blender path: {blender_path}")
360 print(f" Base scene: {base_scene is not None}")
361 if base_scene:
362 print(f" Objects in scene: {len(base_scene.get('objects', []))}")
363 error_messages.append(error_detail)
364 continue
365
366 successful_scenes += 1
367
368 counterfactuals = generate_counterfactuals(
369 base_scene,
370 num_counterfactuals=num_counterfactuals,
371 cf_types=cf_types,
372 same_cf_type=same_cf_type,
373 min_change_score=min_change_score,
374 max_cf_attempts=max_cf_attempts,
375 min_noise_level='light',
376 semantic_only=semantic_only,
377 negative_only=negative_only
378 )
379
380 scene_num = scene_idx + 1
381 scene_prefix = f"scene_{scene_num:04d}"
382
383 base_scene['cf_metadata'] = {
384 'variant': 'original',
385 'is_counterfactual': False,
386 'cf_index': None,
387 'cf_category': 'original',
388 'cf_type': None,
389 'cf_description': None,
390 'source_scene': scene_prefix,
391 }
392 original_scene_path = os.path.join(scenes_dir, f"{scene_prefix}_original.json")
393 save_scene(base_scene, original_scene_path)
394
395 for idx, cf in enumerate(counterfactuals):
396 cf_name = f"cf{idx+1}"
397 cf_scene = cf['scene']
398 cf_scene['cf_metadata'] = {
399 'variant': cf_name,
400 'is_counterfactual': True,
401 'cf_index': idx + 1,
402 'cf_category': cf.get('cf_category', 'unknown'),
403 'cf_type': cf.get('type', None),
404 'cf_description': cf.get('description', None),
405 'change_score': cf.get('change_score', None),
406 'change_attempts': cf.get('change_attempts', None),
407 'source_scene': scene_prefix,
408 }
409 cf_scene_path = os.path.join(scenes_dir, f"{scene_prefix}_{cf_name}.json")
410 save_scene(cf_scene, cf_scene_path)
411
412 render_success = 0
413 total_to_render = len(counterfactuals) + 1
414
415 if not skip_render:
416 if blender_path and blender_available:
417 original_image_path = os.path.join(images_dir, f"{scene_prefix}_original.png")
418 if render_scene(
419 blender_path,
420 original_scene_path,
421 original_image_path,
422 use_gpu=use_gpu,
423 samples=samples,
424 width=width,
425 height=height
426 ):
427 render_success += 1
428
429 for idx, cf in enumerate(counterfactuals):
430 cf_name = f"cf{idx+1}"
431 cf_scene_path = os.path.join(scenes_dir, f"{scene_prefix}_{cf_name}.json")
432 cf_image_path = os.path.join(images_dir, f"{scene_prefix}_{cf_name}.png")
433
434 if render_scene(
435 blender_path,
436 cf_scene_path,
437 cf_image_path,
438 use_gpu=use_gpu,
439 samples=samples,
440 width=width,
441 height=height
442 ):
443 render_success += 1
444
445 if render_success == total_to_render:
446 successful_renders += 1
447 else:
448 print("Blender not available - skipping image rendering. Scene JSON files will still be generated.")
449
450 csv_filename = 'image_mapping_with_questions.csv' if generate_questions else 'image_mapping.csv'
451 csv_path = os.path.join(output_dir, csv_filename)
452
453 try:
454 if generate_mapping_with_questions is not None:
455 generate_mapping_with_questions(
456 run_dir=output_dir,
457 csv_filename=csv_filename,
458 generate_questions=generate_questions,
459 with_links=False,
460 strict_question_validation=True
461 )
462 csv_created = os.path.exists(csv_path)
463 else:
464 csv_created = False
465 except Exception:
466 import traceback
467 traceback.print_exc()
468 csv_created = False
469
470 scene_files = list(Path(scenes_dir).glob("*.json")) if os.path.exists(scenes_dir) else []
471 image_files = list(Path(images_dir).glob("*.png")) if os.path.exists(images_dir) else []
472
473 statistics = {
474 'scenes_generated': successful_scenes,
475 'scenes_rendered': successful_renders,
476 'total_scene_files': len(scene_files),
477 'total_image_files': len(image_files),
478 'num_counterfactuals': num_counterfactuals,
479 'cf_types_used': cf_types if cf_types else 'default',
480 'csv_created': csv_created,
481 'csv_path': csv_path if csv_created else None
482 }
483
484 script_dir = os.path.dirname(os.path.abspath(__file__))
485 cwd = os.getcwd()
486 import shutil
487 import time
488
489 temp_output_dir = os.path.join(cwd, 'temp_output')
490 if os.path.exists(temp_output_dir):
491 for attempt in range(3):
492 try:
493 shutil.rmtree(temp_output_dir)
494 break
495 except Exception as e:
496 if attempt < 2:
497 time.sleep(0.3)
498 else:
499 print(f"Warning: Could not remove temp_output after 3 attempts: {e}")
500
501 render_patched_path = os.path.join(cwd, 'render_images_patched.py')
502 if os.path.exists(render_patched_path):
503 for attempt in range(3):
504 try:
505 time.sleep(0.2)
506 if os.path.exists(render_patched_path):
507 os.remove(render_patched_path)
508 break
509 except Exception as e:
510 if attempt < 2:
511 time.sleep(0.3)
512 else:
513 print(f"Warning: Could not remove render_images_patched.py after 3 attempts: {e}")
514
515 if successful_scenes == 0 and error_messages:
516 error_summary = "Scenes failed. Common reasons:\n"
517 error_summary += "- Blender is not installed or not in PATH\n"
518 error_summary += "- Blender executable not found\n"
519 error_summary += f"\nFirst error: {error_messages[0] if error_messages else 'Unknown error'}"
520
521 return {
522 'success': False,
523 'error': error_summary,
524 'num_scenes': successful_scenes,
525 'output_dir': output_dir,
526 'error_messages': error_messages
527 }
528
529 return {
530 'success': True,
531 'num_scenes': successful_scenes,
532 'output_dir': output_dir,
533 'statistics': statistics,
534 'error_messages': error_messages if error_messages else None
535 }
536
537 except Exception as e:
538 import traceback
539 error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
540 print(error_msg)
541 return {
542 'success': False,
543 'error': error_msg,
544 'num_scenes': successful_scenes,
545 'output_dir': output_dir,
546 'error_messages': error_messages if 'error_messages' in locals() else []
547 }
548
549def main():
550 st.markdown('<p class="main-header">Counterfactual Image Generator</p>', unsafe_allow_html=True)
551
552 if 'output_dir' not in st.session_state:
553 st.session_state.output_dir = None
554
555 if 'generation_complete' not in st.session_state:
556 st.session_state.generation_complete = False
557
558 with st.sidebar:
559 st.header("Configuration")
560
561 st.subheader("Scene Settings")
562 num_scenes = st.number_input(
563 "Number of Scenes",
564 min_value=1,
565 max_value=10000,
566 value=5,
567 help="Number of scene sets to generate"
568 )
569
570 use_fixed_objects = st.checkbox("Use Fixed Number of Objects", value=True)
571
572 if use_fixed_objects:
573 num_objects = st.number_input(
574 "Number of Objects per Scene",
575 min_value=1,
576 max_value=15,
577 value=5,
578 help="Fixed number of objects per scene"
579 )
580 min_objects = None
581 max_objects = None
582 else:
583 num_objects = None
584 min_objects = st.number_input(
585 "Min Objects per Scene",
586 min_value=1,
587 max_value=15,
588 value=3,
589 help="Minimum objects per scene"
590 )
591 max_objects = st.number_input(
592 "Max Objects per Scene",
593 min_value=1,
594 max_value=15,
595 value=7,
596 help="Maximum objects per scene"
597 )
598 if min_objects > max_objects:
599 st.error("Min objects must be <= Max objects")
600 return
601
602 st.subheader("Counterfactual Settings")
603 num_counterfactuals = st.number_input(
604 "Number of Counterfactuals",
605 min_value=1,
606 max_value=10,
607 value=2,
608 help="Number of counterfactual variants per scene"
609 )
610
611 st.markdown("**Counterfactual Types**")
612 st.caption("Leave all unchecked to use default behavior (1 Image CF + 1 Negative CF)")
613 semantic_only = st.checkbox(
614 "Semantic only",
615 value=False,
616 help="Generate only Semantic/Image counterfactuals (Change Color, Add Object, etc.); no Negative CFs"
617 )
618 negative_only = st.checkbox(
619 "Negative only",
620 value=False,
621 help="Generate only Negative counterfactuals (Change Lighting, Add Noise, Occlusion Change, etc.); no Semantic CFs"
622 )
623 same_cf_type = st.checkbox(
624 "Same CF type for all",
625 value=False,
626 help="Use the same counterfactual type for every variant (first selected type, or one random if none selected)"
627 )
628 with st.expander("Image CFs (change answers)", expanded=True):
629 use_change_color = st.checkbox("Change Color", value=False)
630 use_change_shape = st.checkbox("Change Shape", value=False)
631 use_change_size = st.checkbox("Change Size", value=False)
632 use_change_material = st.checkbox("Change Material", value=False)
633 use_change_position = st.checkbox("Change Position", value=False)
634 use_add_object = st.checkbox("Add Object", value=False)
635 use_remove_object = st.checkbox("Remove Object", value=False)
636 use_replace_object = st.checkbox("Replace Object", value=False)
637 use_swap_attribute = st.checkbox("Swap Attribute", value=False)
638 use_relational_flip = st.checkbox("Relational Flip", value=False)
639
640 with st.expander("Negative CFs (don't change answers)", expanded=False):
641 use_change_background = st.checkbox("Change Background", value=False)
642 use_change_lighting = st.checkbox("Change Lighting", value=False)
643 use_add_noise = st.checkbox("Add Noise", value=False)
644 use_occlusion_change = st.checkbox("Occlusion Change", value=False)
645 use_apply_fisheye = st.checkbox("Apply Fisheye", value=False)
646 use_apply_blur = st.checkbox("Apply Blur", value=False)
647 use_apply_vignette = st.checkbox("Apply Vignette", value=False)
648 use_apply_chromatic_aberration = st.checkbox("Apply Chromatic Aberration", value=False)
649
650 with st.expander("Advanced Settings", expanded=False):
651 min_change_score = st.slider(
652 "Minimum Change Score",
653 min_value=0.5,
654 max_value=5.0,
655 value=1.0,
656 step=0.1,
657 help="Minimum heuristic change score for counterfactuals"
658 )
659
660 max_cf_attempts = st.number_input(
661 "Max CF Attempts",
662 min_value=1,
663 max_value=50,
664 value=10,
665 help="Maximum retries per counterfactual"
666 )
667
668 min_noise_level = st.selectbox(
669 "Min Noise Level (for add_noise CF)",
670 options=['light', 'medium', 'heavy'],
671 index=0,
672 help="Minimum noise level when using add_noise counterfactual"
673 )
674
675 st.markdown("---")
676 st.markdown("**Rendering Settings**")
677
678 use_gpu = st.checkbox("Use GPU Rendering", value=False)
679 use_gpu_int = 1 if use_gpu else 0
680
681 samples = st.number_input(
682 "Render Samples",
683 min_value=64,
684 max_value=2048,
685 value=512,
686 step=64,
687 help="Cycles sampling rate (higher = better quality, slower)"
688 )
689
690 image_width = st.number_input(
691 "Image Width",
692 min_value=160,
693 max_value=1920,
694 value=320,
695 step=80
696 )
697
698 image_height = st.number_input(
699 "Image Height",
700 min_value=120,
701 max_value=1080,
702 value=240,
703 step=60
704 )
705
706 st.markdown("**CSV Options**")
707 generate_questions = st.checkbox(
708 "Generate Questions in CSV",
709 value=False,
710 help="Include question and answer columns in the CSV file"
711 )
712
713 cf_types = []
714 if use_change_color:
715 cf_types.append('change_color')
716 if use_change_shape:
717 cf_types.append('change_shape')
718 if use_change_size:
719 cf_types.append('change_size')
720 if use_change_material:
721 cf_types.append('change_material')
722 if use_change_position:
723 cf_types.append('change_position')
724 if use_add_object:
725 cf_types.append('add_object')
726 if use_remove_object:
727 cf_types.append('remove_object')
728 if use_replace_object:
729 cf_types.append('replace_object')
730 if use_swap_attribute:
731 cf_types.append('swap_attribute')
732 if use_relational_flip:
733 cf_types.append('relational_flip')
734 if use_change_background:
735 cf_types.append('change_background')
736 if use_change_lighting:
737 cf_types.append('change_lighting')
738 if use_add_noise:
739 cf_types.append('add_noise')
740 if use_occlusion_change:
741 cf_types.append('occlusion_change')
742 if use_apply_fisheye:
743 cf_types.append('apply_fisheye')
744 if use_apply_blur:
745 cf_types.append('apply_blur')
746 if use_apply_vignette:
747 cf_types.append('apply_vignette')
748 if use_apply_chromatic_aberration:
749 cf_types.append('apply_chromatic_aberration')
750
751 if not cf_types:
752 cf_types = None
753
754 col1, col2 = st.columns([2, 1])
755
756 with col1:
757 st.header("Generate Counterfactual Images")
758
759 if st.button("Generate Counterfactual", use_container_width=True, key="generate_button"):
760 st.session_state.generation_complete = False
761 st.session_state.generating = True
762
763 if num_scenes < 1:
764 st.error("Please specify at least 1 scene to generate.")
765 return
766
767 if use_fixed_objects and num_objects < 1:
768 st.error("Please specify at least 1 object per scene.")
769 return
770 if not use_fixed_objects and (min_objects < 1 or max_objects < 1 or min_objects > max_objects):
771 st.error("Invalid min/max objects configuration.")
772 return
773
774 if os.path.exists('/tmp'):
775 base_dir = '/tmp'
776 else:
777 base_dir = tempfile.gettempdir()
778
779 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
780 output_dir = os.path.join(base_dir, f"counterfactual_output_{timestamp}")
781 os.makedirs(output_dir, exist_ok=True)
782 st.session_state.output_dir = output_dir
783
784 import shutil
785 import time
786 script_dir = os.path.dirname(os.path.abspath(__file__))
787 cwd = os.getcwd()
788
789 temp_output_dir = os.path.join(cwd, 'temp_output')
790 if os.path.exists(temp_output_dir):
791 for attempt in range(3):
792 try:
793 shutil.rmtree(temp_output_dir)
794 break
795 except Exception as e:
796 if attempt < 2:
797 time.sleep(0.3)
798 else:
799 print(f"Warning: Could not remove temp_output after 3 attempts: {e}")
800
801 render_patched_path = os.path.join(cwd, 'render_images_patched.py')
802 if os.path.exists(render_patched_path):
803 for attempt in range(3):
804 try:
805 time.sleep(0.2)
806 if os.path.exists(render_patched_path):
807 os.remove(render_patched_path)
808 break
809 except Exception as e:
810 if attempt < 2:
811 time.sleep(0.3)
812 else:
813 print(f"Warning: Could not remove render_images_patched.py after 3 attempts: {e}")
814
815 try:
816 from pipeline import create_patched_render_script
817 create_patched_render_script()
818 except Exception as e:
819 st.warning(f"Could not create patched render script: {e}")
820
821 params = {
822 'num_scenes': num_scenes,
823 'num_objects': num_objects,
824 'num_counterfactuals': num_counterfactuals,
825 'cf_types': cf_types if cf_types else None,
826 'same_cf_type': same_cf_type,
827 'min_change_score': min_change_score,
828 'max_cf_attempts': max_cf_attempts,
829 'width': image_width,
830 'height': image_height,
831 'output_dir': output_dir
832 }
833
834 progress_bar = st.progress(0)
835 status_text = st.empty()
836
837 try:
838 if not PIPELINE_AVAILABLE:
839 st.error("Pipeline functions are not available. Please check your installation.")
840 return
841
842 status_text.text("Initializing generator...")
843 progress_bar.progress(10)
844
845 if use_fixed_objects:
846 status_text.text(f"Generating {num_scenes} scenes with {num_objects} objects each...")
847 else:
848 status_text.text(f"Generating {num_scenes} scenes with {min_objects}-{max_objects} objects each...")
849 progress_bar.progress(30)
850
851 result = generate_counterfactual_scenes(
852 num_scenes=num_scenes,
853 num_objects=num_objects,
854 min_objects=min_objects,
855 max_objects=max_objects,
856 num_counterfactuals=num_counterfactuals,
857 cf_types=cf_types,
858 same_cf_type=same_cf_type,
859 min_change_score=min_change_score,
860 max_cf_attempts=max_cf_attempts,
861 min_noise_level=min_noise_level,
862 output_dir=output_dir,
863 use_gpu=use_gpu_int,
864 samples=samples,
865 width=image_width,
866 height=image_height,
867 skip_render=False,
868 generate_questions=generate_questions,
869 semantic_only=semantic_only,
870 negative_only=negative_only
871 )
872
873 progress_bar.progress(80)
874 status_text.text("Preparing output...")
875
876 if result and result.get('success', False):
877 num_scenes_generated = result.get('num_scenes', 0)
878
879 if num_scenes_generated == 0:
880 st.warning("No scenes were created. Blender is required and is not available in this environment.")
881 st.info("**To use this application:**\n"
882 "1. Run it locally with Blender installed\n"
883 "2. Use the command-line `pipeline.py` script\n"
884 "3. Install Blender and ensure it's in your system PATH")
885 st.session_state.generation_complete = False
886 else:
887 st.session_state.generation_complete = True
888 progress_bar.progress(100)
889 status_text.text("Done.")
890
891 st.success(f"Successfully generated {num_scenes_generated} scene sets!")
892 st.info(f"Output directory: {output_dir}")
893
894 if 'statistics' in result and result['statistics'].get('csv_created'):
895 csv_path = result['statistics'].get('csv_path')
896 if csv_path:
897 st.success(f"CSV file created: `{os.path.basename(csv_path)}`")
898
899 if 'statistics' in result:
900 stats = result['statistics']
901 st.json(stats)
902 else:
903 error_msg = result.get('error', 'Unknown error occurred') if result else 'Failed'
904 st.error(f"Generation failed: {error_msg}")
905
906 if 'blender' in error_msg.lower() or 'Blender' in error_msg or result.get('num_scenes', 0) == 0:
907 st.warning("**Important:** This application requires Blender to generate scenes. Blender is not available on Hugging Face Spaces.")
908 st.info("**To use this application:**\n"
909 "1. Run it locally with Blender installed\n"
910 "2. Use the command-line `pipeline.py` script\n"
911 "3. Install Blender and ensure it's in your system PATH")
912
913 st.session_state.generation_complete = False
914 st.session_state.generating = False
915
916 except Exception as e:
917 st.error(f"Error during generation: {str(e)}")
918 st.exception(e)
919 st.session_state.generation_complete = False
920 st.session_state.generating = False
921 progress_bar.progress(0)
922 status_text.text("Failed")
923
924 with col2:
925 st.header("Output")
926
927 if st.session_state.generation_complete and st.session_state.output_dir:
928 output_dir = st.session_state.output_dir
929
930 if os.path.exists(output_dir):
931 images_dir = os.path.join(output_dir, 'images')
932 scenes_dir = os.path.join(output_dir, 'scenes')
933
934 scene_files = list(Path(scenes_dir).glob("*.json")) if os.path.exists(scenes_dir) else []
935 image_files = list(Path(images_dir).glob("*.png")) if os.path.exists(images_dir) else []
936 csv_files = list(Path(output_dir).rglob("*.csv"))
937
938 st.success("Complete!")
939 st.metric("Scene Files", len(scene_files))
940 st.metric("CSV Files", len(csv_files))
941 st.metric("Image Files", len(image_files))
942
943 if image_files:
944 st.markdown("---")
945 st.subheader("Generated Images")
946
947 def get_counterfactual_type_from_scene(scene_file):
948 try:
949 with open(scene_file, 'r') as f:
950 scene_data = json.load(f)
951 cf_metadata = scene_data.get('cf_metadata', {})
952 cf_type = cf_metadata.get('cf_type', '')
953 if cf_type:
954 return cf_type.replace('_', ' ').title()
955 except Exception as e:
956 pass
957 return "Counterfactual"
958
959 scene_sets = {}
960 for img_file in image_files:
961 filename = img_file.name
962 if filename.startswith('scene_'):
963 parts = filename.replace('.png', '').split('_')
964 if len(parts) >= 3:
965 scene_num = parts[1]
966 scene_type = parts[2]
967
968 if scene_num not in scene_sets:
969 scene_sets[scene_num] = {}
970
971 scene_sets[scene_num][scene_type] = {
972 'image_path': str(img_file),
973 'filename': filename
974 }
975
976 sorted_scenes = sorted(scene_sets.keys())[:3]
977
978 for scene_idx, scene_num in enumerate(sorted_scenes):
979 scene_data = scene_sets[scene_num]
980
981 if 'original' not in scene_data:
982 continue
983
984 st.markdown(f"### Scene {scene_num}")
985
986 cols = st.columns(3)
987
988 with cols[0]:
989 original = scene_data['original']
990 st.image(original['image_path'], use_container_width=True, caption="Original")
991
992 cf_count = 0
993 for cf_key in ['cf1', 'cf2']:
994 if cf_key in scene_data and cf_count < 2:
995 cf_data = scene_data[cf_key]
996 cf_scene_file = os.path.join(scenes_dir, cf_data['filename'].replace('.png', '.json'))
997 cf_type = get_counterfactual_type_from_scene(cf_scene_file) if os.path.exists(cf_scene_file) else f"Counterfactual {cf_count + 1}"
998
999 with cols[cf_count + 1]:
1000 st.image(cf_data['image_path'], use_container_width=True, caption=cf_type)
1001
1002 cf_count += 1
1003
1004 if scene_idx < len(sorted_scenes) - 1:
1005 st.markdown("---")
1006
1007 st.markdown("---")
1008 st.subheader("Download Output")
1009
1010 zip_filename = f"counterfactual_output_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
1011 zip_path = os.path.join(tempfile.gettempdir(), zip_filename)
1012
1013 try:
1014 create_zip_file(output_dir, zip_path)
1015
1016 file_size = os.path.getsize(zip_path) / (1024 * 1024)
1017
1018 with open(zip_path, 'rb') as f:
1019 st.download_button(
1020 label=f"Download as ZIP ({file_size:.2f} MB)",
1021 data=f.read(),
1022 file_name=zip_filename,
1023 mime="application/zip",
1024 use_container_width=True
1025 )
1026
1027 with st.expander("Output Structure"):
1028 st.text(f"Output directory: {output_dir}")
1029 if scene_files:
1030 st.text(f"\nScene files: {len(scene_files)}")
1031 st.text("Sample files:")
1032 for f in scene_files[:5]:
1033 st.text(f" - {f.name}")
1034 if csv_files:
1035 st.text(f"\nCSV files: {len(csv_files)}")
1036 for f in csv_files:
1037 st.text(f" - {f.name}")
1038 if image_files:
1039 st.text(f"\nImage files: {len(image_files)}")
1040 st.text("Sample files:")
1041 for f in image_files[:5]:
1042 st.text(f" - {f.name}")
1043
1044 except Exception as e:
1045 st.error(f"Error creating zip file: {str(e)}")
1046 else:
1047 st.warning("Output directory not found.")
1048 else:
1049 st.info("Configure parameters and click 'Generate Counterfactual' to start.")
1050
1051 st.markdown("---")
1052 st.markdown(
1053 "<div style='text-align: center; color: #666; padding: 1rem;'>"
1054 "Counterfactual Image Tool | Built with Streamlit"
1055 "</div>",
1056 unsafe_allow_html=True
1057 )
1058
1059if __name__ == "__main__":
1060 main()
1061
1062