"""Tab 2: Images.
Upload and manage fire damage images for AI analysis.
MVP Simplification: Single room - images auto-assigned to the room.
"""
import uuid
import gradio as gr
from typing import Any
from PIL import Image
import io
from ui.state import SessionState, ImageFormData
from ui.components import image_store
from config.settings import settings
def create_tab() -> dict[str, Any]:
"""Create Tab 2 UI components.
Returns:
Dictionary of component references for event wiring.
"""
gr.Markdown("### Fire Damage Images")
gr.Markdown(
f"*Upload up to {settings.max_images_per_assessment} images for AI analysis.*"
)
with gr.Row():
with gr.Column(scale=2):
image_upload = gr.Files(
label="Upload Images (select multiple)",
file_count="multiple",
file_types=["image"],
elem_id="image_upload",
)
image_description = gr.Textbox(
label="Description (optional)",
placeholder="e.g., View of ceiling deck from center aisle",
elem_id="image_description",
info="Applied to all images in batch",
)
with gr.Row():
add_image_btn = gr.Button("Add Images", variant="primary")
clear_upload_btn = gr.Button("Clear", variant="secondary")
with gr.Column(scale=3):
images_gallery = gr.Gallery(
label="Images Added",
columns=3,
height="auto",
elem_id="images_gallery",
)
with gr.Row():
remove_last_btn = gr.Button("Remove Last Image", variant="secondary")
clear_all_btn = gr.Button("Clear All Images", variant="stop")
# Image count and status
with gr.Row():
image_count = gr.Textbox(
label="Images Added",
value="0 / 20",
interactive=False,
)
# Validation status
with gr.Row():
validation_status = gr.HTML(
value="",
elem_id="tab2_validation",
)
# Resume warning (shown when images need re-upload)
with gr.Row():
resume_warning = gr.HTML(
value="",
elem_id="resume_warning",
visible=False,
)
with gr.Row():
back_btn = gr.Button("← Back to Room")
validate_btn = gr.Button(
"Validate & Continue to Observations →",
variant="primary",
)
return {
"image_upload": image_upload,
"image_description": image_description,
"add_image_btn": add_image_btn,
"clear_upload_btn": clear_upload_btn,
"images_gallery": images_gallery,
"remove_last_btn": remove_last_btn,
"clear_all_btn": clear_all_btn,
"image_count": image_count,
"validation_status": validation_status,
"resume_warning": resume_warning,
"back_btn": back_btn,
"validate_btn": validate_btn,
}
def add_image(
session: SessionState,
files: list | None,
description: str,
) -> tuple[SessionState, list[tuple], str, str, None, str]:
"""Add one or more images to the session (batch upload).
Images are automatically associated with the single room.
Args:
session: Current session state.
files: List of uploaded file objects from gr.Files, each with a `name` attribute.
description: Optional description applied to all images.
Returns:
Tuple of (session, gallery_data, validation_html, image_count,
cleared_files, cleared_description).
"""
validation_html = ""
# Validate input
errors = []
if not files or len(files) == 0:
errors.append("Please upload at least one image")
# Check capacity
current_count = len(session.images)
max_allowed = settings.max_images_per_assessment
if files and current_count + len(files) > max_allowed:
remaining = max_allowed - current_count
if remaining <= 0:
errors.append(f"Maximum of {max_allowed} images allowed (already at limit)")
else:
errors.append(f"Can only add {remaining} more image(s) (limit: {max_allowed})")
if errors:
error_items = "".join(f"{e}" for e in errors)
validation_html = f"""
"""
gallery_data = _get_gallery_data(session)
count_str = f"{len(session.images)} / {max_allowed}"
return session, gallery_data, validation_html, count_str, files, description
# Get room info from session (single room)
room_id = session.room.id
room_name = session.room.name.replace(" ", "_")[:20] if session.room.name else "room"
# Process each uploaded file
added_count = 0
for file_obj in files:
# Check if we've hit the limit
if len(session.images) >= max_allowed:
break
try:
# Open image from file path
img = Image.open(file_obj.name)
# Generate image ID
image_id = f"img-{uuid.uuid4().hex[:8]}"
# Store image bytes in memory
img_bytes = io.BytesIO()
img.save(img_bytes, format="PNG")
image_store.store(image_id, img_bytes.getvalue())
# Add image metadata to session
img_meta = ImageFormData(
id=image_id,
filename=f"{room_name}_{image_id}.png",
room_id=room_id,
description=description.strip() if description else "",
)
session.images.append(img_meta)
added_count += 1
except Exception:
# Skip files that can't be opened as images
continue
session.update_timestamp()
# Success message
if added_count > 0:
validation_html = f"""
✓ {added_count} image(s) added for {room_name}
"""
else:
validation_html = """
No images could be processed
"""
gallery_data = _get_gallery_data(session)
count_str = f"{len(session.images)} / {max_allowed}"
# Clear form
return session, gallery_data, validation_html, count_str, None, ""
def remove_last_image(session: SessionState) -> tuple[SessionState, list[tuple], str, str]:
"""Remove the last image from the session."""
validation_html = ""
if session.images:
removed = session.images.pop()
image_store.remove(removed.id)
session.update_timestamp()
validation_html = f"""
Removed image: {removed.filename}
"""
gallery_data = _get_gallery_data(session)
count_str = f"{len(session.images)} / {settings.max_images_per_assessment}"
return session, gallery_data, validation_html, count_str
def clear_all_images(session: SessionState) -> tuple[SessionState, list[tuple], str, str]:
"""Clear all images from the session."""
count = len(session.images)
# Clear from store
for img in session.images:
image_store.remove(img.id)
session.images = []
session.update_timestamp()
validation_html = ""
if count > 0:
validation_html = f"""
Cleared {count} image(s)
"""
count_str = f"0 / {settings.max_images_per_assessment}"
return session, [], validation_html, count_str
def validate_and_continue(session: SessionState) -> tuple[SessionState, str, int]:
"""Validate Tab 2 and proceed to Tab 3.
Returns:
Tuple of (session, validation_html, next_tab_index).
"""
# Check if images need re-upload (session restored but images not in memory)
expected_ids = [img.id for img in session.images]
missing_ids = image_store.get_missing_ids(expected_ids)
if missing_ids:
missing_count = len(missing_ids)
html = f"""
⚠ {missing_count} image(s) need to be re-uploaded
Images are not stored in browser storage. Please re-upload the missing images
or clear the image list and start fresh.
"""
return session, html, gr.update(selected=1) # Stay on Images tab (index 1)
is_valid, errors = session.validate_tab2()
if is_valid:
session.tab2_complete = True
session.update_timestamp()
html = """
✓ Images complete. Proceeding to Observations tab...
"""
return session, html, gr.update(selected=2) # Go to tab index 2 (Observations)
else:
session.tab2_complete = False
error_items = "".join(f"{e}" for e in errors)
html = f"""
Please fix the following:
"""
return session, html, gr.update(selected=1) # Stay on current tab
def load_from_session(session: SessionState) -> tuple[list[tuple], str, str]:
"""Load gallery data and count from session.
Returns:
Tuple of (gallery_data, image_count, resume_warning_html).
"""
gallery_data = _get_gallery_data(session)
count_str = f"{len(session.images)} / {settings.max_images_per_assessment}"
# Check for missing images
expected_ids = [img.id for img in session.images]
missing_ids = image_store.get_missing_ids(expected_ids)
resume_html = ""
if missing_ids and session.images:
resume_html = f"""
⚠ {len(missing_ids)} image(s) need to be re-uploaded
Session restored, but images must be re-uploaded as they are not stored in browser storage.
"""
return gallery_data, count_str, resume_html
def _get_gallery_data(session: SessionState) -> list[tuple]:
"""Get gallery data from session images.
Returns:
List of (image, caption) tuples for gallery.
"""
gallery_data = []
for img_meta in session.images:
img_bytes = image_store.get(img_meta.id)
if img_bytes:
# Convert bytes to PIL Image for gallery
pil_image = Image.open(io.BytesIO(img_bytes))
caption = img_meta.description or img_meta.filename
gallery_data.append((pil_image, caption))
return gallery_data